Update tests

pamelafox · pamelafox · commit c124a843e51f · 2024-08-03T16:41:12.000Z
diff --git a/.vscode/settings.json b/.vscode/settings.json
@@ -22,5 +22,10 @@
                 "ssl": true
             }
         }
-    ]
+    ],
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
 }
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -53,7 +53,6 @@ def mock_session_env(monkeypatch_session):
         monkeypatch_session.setenv("POSTGRES_DATABASE", POSTGRES_DATABASE)
         monkeypatch_session.setenv("POSTGRES_PASSWORD", POSTGRES_PASSWORD)
         monkeypatch_session.setenv("POSTGRES_SSL", POSTGRES_SSL)
-        monkeypatch_session.setenv("POSTGRESQL_DATABASE_URL", POSTGRESQL_DATABASE_URL)
         monkeypatch_session.setenv("RUNNING_IN_PRODUCTION", "False")
         # Azure Subscription
         monkeypatch_session.setenv("AZURE_SUBSCRIPTION_ID", "test-storage-subid")
@@ -68,11 +67,62 @@ def mock_session_env(monkeypatch_session):
         monkeypatch_session.setenv("AZURE_OPENAI_EMBED_MODEL", "text-embedding-ada-002")
         monkeypatch_session.setenv("AZURE_OPENAI_EMBED_MODEL_DIMENSIONS", "1536")
         monkeypatch_session.setenv("AZURE_OPENAI_KEY", "fakekey")
-        # Allowed Origin
-        monkeypatch_session.setenv("ALLOWED_ORIGIN", "https://frontend.com")
 
-        if os.getenv("AZURE_USE_AUTHENTICATION") is not None:
-            monkeypatch_session.delenv("AZURE_USE_AUTHENTICATION")
+        yield
+
+
+@pytest.fixture(scope="session")
+def mock_session_env_openai(monkeypatch_session):
+    """Mock the environment variables for testing."""
+    # Note that this does *not* clear existing env variables by default-
+    # we used to specify clear=True but this caused issues with Playwright tests
+    # https://github.com/microsoft/playwright-python/issues/2506
+    with mock.patch.dict(os.environ):
+        # Database
+        monkeypatch_session.setenv("POSTGRES_HOST", POSTGRES_HOST)
+        monkeypatch_session.setenv("POSTGRES_USERNAME", POSTGRES_USERNAME)
+        monkeypatch_session.setenv("POSTGRES_DATABASE", POSTGRES_DATABASE)
+        monkeypatch_session.setenv("POSTGRES_PASSWORD", POSTGRES_PASSWORD)
+        monkeypatch_session.setenv("POSTGRES_SSL", POSTGRES_SSL)
+        monkeypatch_session.setenv("RUNNING_IN_PRODUCTION", "False")
+        # Azure Subscription
+        monkeypatch_session.setenv("AZURE_SUBSCRIPTION_ID", "test-storage-subid")
+        # OpenAI.com OpenAI
+        monkeypatch_session.setenv("OPENAI_CHAT_HOST", "openai")
+        monkeypatch_session.setenv("OPENAI_EMBED_HOST", "openai")
+        monkeypatch_session.setenv("OPENAICOM_KEY", "fakekey")
+        monkeypatch_session.setenv("OPENAICOM_CHAT_MODEL", "gpt-3.5-turbo")
+        monkeypatch_session.setenv("OPENAICOM_EMBED_MODEL", "text-embedding-ada-002")
+        monkeypatch_session.setenv("OPENAICOM_EMBED_MODEL_DIMENSIONS", "1536")
+        monkeypatch_session.setenv("OPENAICOM_EMBEDDING_COLUMN", "embedding_ada002")
+
+        yield
+
+
+@pytest.fixture(scope="session")
+def mock_session_env_ollama(monkeypatch_session):
+    """Mock the environment variables for testing."""
+    # Note that this does *not* clear existing env variables by default-
+    # we used to specify clear=True but this caused issues with Playwright tests
+    # https://github.com/microsoft/playwright-python/issues/2506
+    with mock.patch.dict(os.environ):
+        # Database
+        monkeypatch_session.setenv("POSTGRES_HOST", POSTGRES_HOST)
+        monkeypatch_session.setenv("POSTGRES_USERNAME", POSTGRES_USERNAME)
+        monkeypatch_session.setenv("POSTGRES_DATABASE", POSTGRES_DATABASE)
+        monkeypatch_session.setenv("POSTGRES_PASSWORD", POSTGRES_PASSWORD)
+        monkeypatch_session.setenv("POSTGRES_SSL", POSTGRES_SSL)
+        monkeypatch_session.setenv("RUNNING_IN_PRODUCTION", "False")
+        # Azure Subscription
+        monkeypatch_session.setenv("AZURE_SUBSCRIPTION_ID", "test-storage-subid")
+        # Ollama OpenAI
+        monkeypatch_session.setenv("OPENAI_CHAT_HOST", "ollama")
+        monkeypatch_session.setenv("OPENAI_EMBED_HOST", "ollama")
+        monkeypatch_session.setenv("OLLAMA_ENDPOINT", "http://host.docker.internal:11434/v1")
+        monkeypatch_session.setenv("OLLAMA_CHAT_MODEL", "llama3.1")
+        monkeypatch_session.setenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
+        monkeypatch_session.setenv("OLLAMA_EMBEDDING_COLUMN", "embedding_nomic")
+
         yield
 
 
diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py
@@ -13,6 +13,26 @@ async def test_get_common_parameters(mock_session_env):
     assert result.openai_embed_deployment == "text-embedding-ada-002"
 
 
+@pytest.mark.asyncio
+async def test_get_common_parameters_ollama(mock_session_env_ollama):
+    result = await common_parameters()
+    assert result.openai_chat_model == "llama3.1"
+    assert result.openai_embed_model == "nomic-embed-text"
+    assert result.openai_embed_dimensions is None
+    assert result.openai_chat_deployment is None
+    assert result.openai_embed_deployment is None
+
+
+@pytest.mark.asyncio
+async def test_get_common_parameters_openai(mock_session_env_openai):
+    result = await common_parameters()
+    assert result.openai_chat_model == "gpt-3.5-turbo"
+    assert result.openai_embed_model == "text-embedding-ada-002"
+    assert result.openai_embed_dimensions == 1536
+    assert result.openai_chat_deployment is None
+    assert result.openai_embed_deployment is None
+
+
 @pytest.mark.asyncio
 async def test_get_azure_credentials(mock_session_env, mock_default_azure_credential):
     result = await get_azure_credentials()
diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py
@@ -13,6 +13,18 @@ async def test_compute_text_embedding(mock_default_azure_credential, mock_openai
         openai_client=openai_embed_client,
         embed_model="text-embedding-ada-002",
         embed_deployment="text-embedding-ada-002",
+    )
+    assert result == test_data.embeddings
+
+
+@pytest.mark.asyncio
+async def test_compute_text_embedding_dimensions(mock_default_azure_credential, mock_openai_embedding):
+    openai_embed_client = await create_openai_embed_client(mock_default_azure_credential)
+    result = await compute_text_embedding(
+        q="test",
+        openai_client=openai_embed_client,
+        embed_model="text-embedding-3-small",
+        embed_deployment="text-embedding-3-small",
         embedding_dimensions=1536,
     )
     assert result == test_data.embeddings

Original file line number	Diff line number	Diff line change
`@@ -22,5 +22,10 @@`
`22`	`22`	`"ssl": true`
`23`	`23`	`}`
`24`	`24`	`}`
`25`		`- ]`
	`25`	`+ ],`
	`26`	`+ "python.testing.pytestArgs": [`
	`27`	`+ "tests"`
	`28`	`+ ],`
	`29`	`+ "python.testing.unittestEnabled": false,`
	`30`	`+ "python.testing.pytestEnabled": true`
`26`	`31`	`}`