add more tests and fix azure credentials mocking

AsyncTokenCredentials is not the correct class to inherit from as we are not using the async credentials

Author: john0isaac

src/fastapi_app/embeddings.py

@@ -11,7 +11,7 @@ async def compute_text_embedding(
     embed_model: str,
     embed_deployment: str | None = None,
     embedding_dimensions: int = 1536,
-):
+) -> list[float]:
     SUPPORTED_DIMENSIONS_MODEL = {
         "text-embedding-ada-002": False,
         "text-embedding-3-small": True,

tests/conftest.py

@@ -224,7 +224,7 @@ def mock_default_azure_credential(mock_session_env):
     """Mock the Azure credential for testing."""
     with mock.patch("azure.identity.DefaultAzureCredential") as mock_default_azure_credential:
         mock_default_azure_credential.return_value = MockAzureCredential()
-        yield
+        yield mock_default_azure_credential
 
 
 @pytest_asyncio.fixture(scope="function")

tests/mocks.py

@@ -1,34 +1,22 @@
 from collections import namedtuple
-from types import TracebackType
 
-from azure.core.credentials_async import AsyncTokenCredential
+from azure.core.credentials import TokenCredential
 
-MockToken = namedtuple("MockToken", ["token", "expires_on", "value"])
+MockToken = namedtuple("MockToken", ["token", "expires_on"])
 
 
-class MockAzureCredential(AsyncTokenCredential):
-    async def get_token(self, uri):
-        return MockToken("", 9999999999, "")
-
-    async def close(self) -> None:
-        pass
-
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None = None,
-        exc_value: BaseException | None = None,
-        traceback: TracebackType | None = None,
-    ) -> None:
-        pass
+class MockAzureCredential(TokenCredential):
+    def get_token(self, uri):
+        return MockToken("", 9999999999)
 
 
-class MockAzureCredentialExpired(AsyncTokenCredential):
+class MockAzureCredentialExpired(TokenCredential):
     def __init__(self):
         self.access_number = 0
 
     async def get_token(self, uri):
         self.access_number += 1
         if self.access_number == 1:
-            return MockToken("", 0, "")
+            return MockToken("", 0)
         else:
-            return MockToken("", 9999999999, "")
+            return MockToken("", 9999999999)

tests/test_dependencies.py

@@ -0,0 +1,21 @@
+import pytest
+
+from fastapi_app.dependencies import common_parameters, get_azure_credentials
+
+
+@pytest.mark.asyncio
+async def test_get_common_parameters(mock_session_env):
+    result = await common_parameters()
+    assert result.openai_chat_model == "gpt-35-turbo"
+    assert result.openai_embed_model == "text-embedding-ada-002"
+    assert result.openai_embed_dimensions == 1536
+    assert result.openai_chat_deployment == "gpt-35-turbo"
+    assert result.openai_embed_deployment == "text-embedding-ada-002"
+
+
+@pytest.mark.asyncio
+async def test_get_azure_credentials(mock_session_env, mock_default_azure_credential):
+    result = await get_azure_credentials()
+    token = result.get_token("https://vault.azure.net")
+    assert token.expires_on == 9999999999
+    assert token.token == ""

tests/test_embeddings.py

@@ -0,0 +1,18 @@
+import pytest
+
+from fastapi_app.embeddings import compute_text_embedding
+from fastapi_app.openai_clients import create_openai_embed_client
+from tests.data import test_data
+
+
+@pytest.mark.asyncio
+async def test_compute_text_embedding(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-ada-002",
+        embed_deployment="text-embedding-ada-002",
+        embedding_dimensions=1536,
+    )
+    assert result == test_data.embeddings

tests/test_openai_clients.py

@@ -0,0 +1,24 @@
+import pytest
+
+from fastapi_app.openai_clients import create_openai_chat_client, create_openai_embed_client
+from tests.data import test_data
+
+
+@pytest.mark.asyncio
+async def test_create_openai_embed_client(mock_default_azure_credential, mock_openai_embedding):
+    openai_embed_client = await create_openai_embed_client(mock_default_azure_credential)
+    assert openai_embed_client.embeddings.create is not None
+    embeddings = await openai_embed_client.embeddings.create(
+        model="text-embedding-ada-002", input="test", dimensions=1536
+    )
+    assert embeddings.data[0].embedding == test_data.embeddings
+
+
+@pytest.mark.asyncio
+async def test_create_openai_chat_client(mock_default_azure_credential, mock_openai_chatcompletion):
+    openai_chat_client = await create_openai_chat_client(mock_default_azure_credential)
+    assert openai_chat_client.chat.completions.create is not None
+    response = await openai_chat_client.chat.completions.create(
+        model="gpt-35-turbo", messages=[{"content": "test", "role": "user"}]
+    )
+    assert response.choices[0].message.content == "The capital of France is Paris. [Benefit_Options-2.pdf]."

tests/test_postgres_engine.py

@@ -0,0 +1,63 @@
+import os
+
+import pytest
+
+from fastapi_app.postgres_engine import (
+    create_postgres_engine,
+    create_postgres_engine_from_args,
+    create_postgres_engine_from_env,
+)
+from tests.conftest import POSTGRES_DATABASE, POSTGRES_HOST, POSTGRES_PASSWORD, POSTGRES_SSL, POSTGRES_USERNAME
+
+
+@pytest.mark.asyncio
+async def test_create_postgres_engine(mock_session_env, mock_default_azure_credential):
+    engine = await create_postgres_engine(
+        host=os.environ["POSTGRES_HOST"],
+        username=os.environ["POSTGRES_USERNAME"],
+        database=os.environ["POSTGRES_DATABASE"],
+        password=os.environ.get("POSTGRES_PASSWORD"),
+        sslmode=os.environ.get("POSTGRES_SSL"),
+        azure_credential=mock_default_azure_credential,
+    )
+    assert engine.url.host == "localhost"
+    assert engine.url.username == "admin"
+    assert engine.url.database == "postgres"
+    assert engine.url.password == "postgres"
+    assert engine.url.query["ssl"] == "prefer"
+
+
+@pytest.mark.asyncio
+async def test_create_postgres_engine_from_env(mock_session_env, mock_default_azure_credential):
+    engine = await create_postgres_engine_from_env(
+        azure_credential=mock_default_azure_credential,
+    )
+    assert engine.url.host == "localhost"
+    assert engine.url.username == "admin"
+    assert engine.url.database == "postgres"
+    assert engine.url.password == "postgres"
+    assert engine.url.query["ssl"] == "prefer"
+
+
+@pytest.mark.asyncio
+async def test_create_postgres_engine_from_args(mock_default_azure_credential):
+    args = type(
+        "Args",
+        (),
+        {
+            "host": POSTGRES_HOST,
+            "username": POSTGRES_USERNAME,
+            "database": POSTGRES_DATABASE,
+            "password": POSTGRES_PASSWORD,
+            "sslmode": POSTGRES_SSL,
+        },
+    )
+    engine = await create_postgres_engine_from_args(
+        args=args,
+        azure_credential=mock_default_azure_credential,
+    )
+    assert engine.url.host == "localhost"
+    assert engine.url.username == "admin"
+    assert engine.url.database == "postgres"
+    assert engine.url.password == "postgres"
+    assert engine.url.query["ssl"] == "prefer"