add mocks and use monkey patch for setting env vars

Author: john0isaac

src/fastapi_app/__init__.py

@@ -52,11 +52,11 @@ async def lifespan(app: FastAPI):
     await engine.dispose()
 
 
-def create_app(is_testing: bool = False):
+def create_app(testing: bool = False):
     env = Env()
 
     if not os.getenv("RUNNING_IN_PRODUCTION"):
-        if not is_testing:
+        if not testing:
             env.read_env(".env")
         logging.basicConfig(level=logging.INFO)
     else:

tests/__init__.py

@@ -1,13 +1,14 @@
 import os
 from pathlib import Path
-from unittest.mock import patch
+from unittest import mock
 
 import pytest
 from fastapi.testclient import TestClient
 from sqlalchemy.ext.asyncio import async_sessionmaker
 
 from fastapi_app import create_app
 from fastapi_app.globals import global_storage
+from tests.mocks import MockAzureCredential
 
 POSTGRES_HOST = "localhost"
 POSTGRES_USERNAME = "admin"
@@ -20,34 +21,54 @@
 
 
 @pytest.fixture(scope="session")
-def setup_env():
-    os.environ["POSTGRES_HOST"] = POSTGRES_HOST
-    os.environ["POSTGRES_USERNAME"] = POSTGRES_USERNAME
-    os.environ["POSTGRES_DATABASE"] = POSTGRES_DATABASE
-    os.environ["POSTGRES_PASSWORD"] = POSTGRES_PASSWORD
-    os.environ["POSTGRES_SSL"] = POSTGRES_SSL
-    os.environ["POSTGRESQL_DATABASE_URL"] = POSTGRESQL_DATABASE_URL
-    os.environ["RUNNING_IN_PRODUCTION"] = "False"
-    os.environ["OPENAI_API_KEY"] = "fakekey"
+def monkeypatch_session():
+    with pytest.MonkeyPatch.context() as monkeypatch_session:
+        yield monkeypatch_session
 
 
 @pytest.fixture(scope="session")
-def mock_azure_credential():
-    """Mock the Azure credential for testing."""
-    with patch("azure.identity.DefaultAzureCredential", return_value=None):
+def mock_session_env(monkeypatch_session):
+    """Mock the environment variables for testing."""
+    with mock.patch.dict(os.environ, clear=True):
+        # 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("POSTGRESQL_DATABASE_URL", POSTGRESQL_DATABASE_URL)
+        monkeypatch_session.setenv("RUNNING_IN_PRODUCTION", "False")
+        # Azure Subscription
+        monkeypatch_session.setenv("AZURE_SUBSCRIPTION_ID", "test-storage-subid")
+        # OpenAI
+        monkeypatch_session.setenv("AZURE_OPENAI_CHATGPT_MODEL", "gpt-35-turbo")
+        monkeypatch_session.setenv("OPENAI_API_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 app(setup_env, mock_azure_credential):
+def app(mock_session_env):
     """Create a FastAPI app."""
     if not Path("src/static/").exists():
         pytest.skip("Please generate frontend files first!")
-    return create_app(is_testing=True)
+    return create_app(testing=True)
+
+
+@pytest.fixture(scope="function")
+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 mock_default_azure_credential
 
 
 @pytest.fixture(scope="function")
-def test_client(app):
+def test_client(monkeypatch, app, mock_default_azure_credential):
     """Create a test client."""
 
     with TestClient(app) as test_client:

tests/conftest.py

@@ -0,0 +1,108 @@
+import json
+from collections import namedtuple
+
+import openai.types
+from azure.core.credentials_async import AsyncTokenCredential
+
+MOCK_EMBEDDING_DIMENSIONS = 1536
+MOCK_EMBEDDING_MODEL_NAME = "text-embedding-ada-002"
+
+MockToken = namedtuple("MockToken", ["token", "expires_on", "value"])
+
+
+class MockAzureCredential(AsyncTokenCredential):
+    async def get_token(self, uri):
+        return MockToken("", 9999999999, "")
+
+
+class MockAzureCredentialExpired(AsyncTokenCredential):
+    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, "")
+        else:
+            return MockToken("", 9999999999, "")
+
+
+class MockAsyncPageIterator:
+    def __init__(self, data):
+        self.data = data
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        if not self.data:
+            raise StopAsyncIteration
+        return self.data.pop(0)  # This should be a list of dictionaries.
+
+
+class MockCaption:
+    def __init__(self, text, highlights=None, additional_properties=None):
+        self.text = text
+        self.highlights = highlights or []
+        self.additional_properties = additional_properties or {}
+
+
+class MockResponse:
+    def __init__(self, text, status):
+        self.text = text
+        self.status = status
+
+    async def text(self):
+        return self._text
+
+    async def __aexit__(self, exc_type, exc, tb):
+        pass
+
+    async def __aenter__(self):
+        return self
+
+    async def json(self):
+        return json.loads(self.text)
+
+
+class MockEmbeddingsClient:
+    def __init__(self, create_embedding_response: openai.types.CreateEmbeddingResponse):
+        self.create_embedding_response = create_embedding_response
+
+    async def create(self, *args, **kwargs) -> openai.types.CreateEmbeddingResponse:
+        return self.create_embedding_response
+
+
+class MockClient:
+    def __init__(self, embeddings_client):
+        self.embeddings = embeddings_client
+
+
+def mock_computervision_response():
+    return MockResponse(
+        status=200,
+        text=json.dumps(
+            {
+                "vector": [
+                    0.011925711,
+                    0.023533698,
+                    0.010133852,
+                    0.0063544377,
+                    -0.00038590943,
+                    0.0013952175,
+                    0.009054946,
+                    -0.033573493,
+                    -0.002028305,
+                ],
+                "modelVersion": "2022-04-11",
+            }
+        ),
+    )
+
+
+class MockSynthesisResult:
+    def __init__(self, result):
+        self.__result = result
+
+    def get(self):
+        return self.__result