missed a spot

Signed-off-by: Adrian Cole <adrian.cole@elastic.co>
square · Sep 20, 2024 · 26e48dc · 26e48dc
1 parent 0ea3af4
commit 26e48dc
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/providers/test_ollama.py b/tests/providers/test_ollama.py
@@ -1,3 +1,4 @@
+import os
 import pytest
 from exchange import Message
 from exchange.providers.ollama import OllamaProvider, OLLAMA_MODEL
@@ -6,7 +7,7 @@
 @pytest.mark.integration
 def test_ollama_integration():
     provider = OllamaProvider.from_env()
-    model = OLLAMA_MODEL
+    model = os.getenv("OLLAMA_MODEL", OLLAMA_MODEL)
     system = "You are a helpful assistant."
     messages = [Message.user("Hello")]