Update the tests

Author: pamelafox

app/backend/app.py

@@ -16,7 +16,6 @@
     SpeechSynthesisResult,
     SpeechSynthesizer,
 )
-from azure.core.exceptions import ResourceNotFoundError
 from azure.identity.aio import (
     AzureDeveloperCliCredential,
     ManagedIdentityCredential,
@@ -155,16 +154,12 @@ async def content_file(path: str, auth_claims: dict[str, Any]):
     if blob is None:
         current_app.logger.info("Path not found in general Blob container: %s", path)
         if current_app.config[CONFIG_USER_UPLOAD_ENABLED]:
-            try:
-                user_oid = auth_claims["oid"]
-                user_blob_manager: AdlsBlobManager = current_app.config[CONFIG_USER_BLOB_MANAGER]
-                blob = await user_blob_manager.download_blob(path, user_oid=user_oid)
-            except ResourceNotFoundError:
+            user_oid = auth_claims["oid"]
+            user_blob_manager: AdlsBlobManager = current_app.config[CONFIG_USER_BLOB_MANAGER]
+            blob = await user_blob_manager.download_blob(path, user_oid=user_oid)
+            if blob is None:
                 current_app.logger.exception("Path not found in DataLake: %s", path)
-                abort(404)
-        else:
-            abort(404)
-    if not blob.properties or not blob.properties.has_key("content_settings"):
+    if not blob or not blob.properties or not blob.properties.has_key("content_settings"):
         abort(404)
     mime_type = blob.properties["content_settings"]["content_type"]
     if mime_type == "application/octet-stream":

app/backend/prepdocslib/blobmanager.py

@@ -287,7 +287,7 @@ async def download_blob(
             filename = path_parts[-1]
 
         try:
-            user_directory_client = self._ensure_directory(directory_path=directory_path, user_oid=user_oid)
+            user_directory_client = await self._ensure_directory(directory_path=directory_path, user_oid=user_oid)
             file_client = user_directory_client.get_file_client(filename)
             blob = await file_client.download_file()
             return blob
@@ -343,10 +343,10 @@ async def list_blobs(self, user_oid: str) -> list[str]:
         Returns:
             list[str]: List of filenames that belong to the user
         """
-        user_directory_client = await self._ensure_directory(directory_path=user_oid, user_oid=user_oid)
+        await self._ensure_directory(directory_path=user_oid, user_oid=user_oid)
         files = []
         try:
-            all_paths = user_directory_client.get_paths()
+            all_paths = self.file_system_client.get_paths(path=user_oid, recursive=True)
             async for path in all_paths:
                 # Split path into parts (user_oid/filename or user_oid/directory/files)
                 path_parts = path.name.split("/", 1)

tests/conftest.py

@@ -18,7 +18,7 @@
     SearchField,
     SearchIndex,
 )
-from azure.storage.blob.aio import ContainerClient
+from azure.storage.blob.aio import BlobServiceClient, ContainerClient
 from openai.types import CompletionUsage, CreateEmbeddingResponse, Embedding
 from openai.types.chat import ChatCompletion, ChatCompletionChunk
 from openai.types.chat.chat_completion import (
@@ -43,6 +43,7 @@
     MockAzureCredentialExpired,
     MockBlobClient,
     MockResponse,
+    MockTransport,
     mock_retrieval_response,
     mock_speak_text_cancelled,
     mock_speak_text_failed,
@@ -299,6 +300,14 @@ def mock_blob_container_client(monkeypatch):
     monkeypatch.setattr(ContainerClient, "get_blob_client", lambda *args, **kwargs: MockBlobClient())
 
 
+@pytest.fixture
+def mock_blob_container_client_exists(monkeypatch):
+    async def mock_exists(*args, **kwargs):
+        return True
+
+    monkeypatch.setattr("azure.storage.blob.aio.ContainerClient.exists", mock_exists)
+
+
 envs = [
     {
         "OPENAI_HOST": "openai",
@@ -731,7 +740,7 @@ async def vision_client(
     mock_openai_chatcompletion,
     mock_openai_embedding,
     mock_acs_search,
-    mock_blob_container_client,
+    mock_blob_container_client_exists,
     mock_azurehttp_calls,
 ):
     quart_app = app.create_app()
@@ -740,6 +749,15 @@ async def vision_client(
         test_app.app.config.update({"TESTING": True})
         mock_openai_chatcompletion(test_app.app.config[app.CONFIG_OPENAI_CLIENT])
         mock_openai_embedding(test_app.app.config[app.CONFIG_OPENAI_CLIENT])
+        # need to mock app.CONFIG_GLOBAL_BLOB_MANAGER to
+        mock_blob_service_client = BlobServiceClient(
+            f"https://{os.environ['AZURE_STORAGE_ACCOUNT']}.blob.core.windows.net",
+            credential=MockAzureCredential(),
+            transport=MockTransport(),
+            retry_total=0,  # Necessary to avoid unnecessary network requests during tests
+        )
+        test_app.app.config[app.CONFIG_GLOBAL_BLOB_MANAGER].blob_service_client = mock_blob_service_client
+
         yield test_app.test_client()
 
 
@@ -939,13 +957,14 @@ def mock_get_paths(self, *args, **kwargs):
     monkeypatch.setattr(azure.storage.filedatalake.aio.FileSystemClient, "__init__", mock_init_filesystem_aio)
     monkeypatch.setattr(azure.storage.filedatalake.aio.FileSystemClient, "__aenter__", mock_aenter)
     monkeypatch.setattr(azure.storage.filedatalake.aio.FileSystemClient, "__aexit__", mock_aexit)
-    monkeypatch.setattr(azure.storage.filedatalake.aio.FileSystemClient, "get_paths", mock_get_paths)
+
     monkeypatch.setattr(azure.storage.filedatalake.aio.FileSystemClient, "exists", mock_exists_aio)
-    monkeypatch.setattr(azure.storage.filedatalake.aio.FileSystemClient, "get_paths", mock_get_paths)
     monkeypatch.setattr(azure.storage.filedatalake.aio.FileSystemClient, "get_file_client", mock_get_file_client)
     monkeypatch.setattr(
         azure.storage.filedatalake.aio.FileSystemClient, "create_file_system", mock_create_filesystem_aio
     )
+    monkeypatch.setattr(azure.storage.filedatalake.aio.FileSystemClient, "get_paths", mock_get_paths)
+
     monkeypatch.setattr(azure.storage.filedatalake.aio.FileSystemClient, "create_directory", mock_create_directory_aio)
     monkeypatch.setattr(
         azure.storage.filedatalake.aio.FileSystemClient,
@@ -985,7 +1004,6 @@ async def mock_upload_data_aio(self, *args, **kwargs):
     monkeypatch.setattr(
         azure.storage.filedatalake.aio.DataLakeFileClient, "get_access_control", mock_get_access_control
     )
-
     monkeypatch.setattr(azure.storage.filedatalake.aio.DataLakeFileClient, "__init__", mock_init_file)
     monkeypatch.setattr(azure.storage.filedatalake.aio.DataLakeFileClient, "url", property(mock_url))
     monkeypatch.setattr(azure.storage.filedatalake.aio.DataLakeFileClient, "__aenter__", mock_aenter)
@@ -1000,6 +1018,12 @@ def mock_init_directory(self, path, *args, **kwargs):
         self.path = path
         self.files = {}
 
+    async def mock_get_directory_properties(self, *args, **kwargs):
+        return azure.storage.filedatalake.DirectoryProperties()
+
+    async def mock_get_access_control(self, *args, **kwargs):
+        return {"owner": "OID_X"}
+
     def mock_directory_get_file_client(self, *args, **kwargs):
         path = kwargs.get("file")
         if path in self.files:
@@ -1030,6 +1054,14 @@ async def mock_close_aio(self, *args, **kwargs):
         "update_access_control_recursive",
         mock_update_access_control_recursive_aio,
     )
+    monkeypatch.setattr(
+        azure.storage.filedatalake.aio.DataLakeDirectoryClient,
+        "get_directory_properties",
+        mock_get_directory_properties,
+    )
+    monkeypatch.setattr(
+        azure.storage.filedatalake.aio.DataLakeDirectoryClient, "get_access_control", mock_get_access_control
+    )
     monkeypatch.setattr(azure.storage.filedatalake.aio.DataLakeDirectoryClient, "close", mock_close_aio)
 
     def mock_readinto(self, stream: IO[bytes]):

tests/mocks.py

@@ -3,9 +3,16 @@
 from io import BytesIO
 from typing import Optional
 
+import aiohttp
 import openai.types
 from azure.cognitiveservices.speech import ResultReason
 from azure.core.credentials_async import AsyncTokenCredential
+from azure.core.exceptions import ResourceNotFoundError
+from azure.core.pipeline.transport import (
+    AioHttpTransportResponse,
+    AsyncHttpTransport,
+    HttpRequest,
+)
 from azure.search.documents.agent.models import (
     KnowledgeAgentAzureSearchDocReference,
     KnowledgeAgentMessage,
@@ -63,6 +70,57 @@ async def readinto(self, buffer: BytesIO):
         buffer.write(b"test")
 
 
+class MockAiohttpClientResponse404(aiohttp.ClientResponse):
+    def __init__(self, url, body_bytes, headers=None):
+        self._body = body_bytes
+        self._headers = headers
+        self._cache = {}
+        self.status = 404
+        self.reason = "Not Found"
+        self._url = url
+
+
+class MockAiohttpClientResponse(aiohttp.ClientResponse):
+    def __init__(self, url, body_bytes, headers=None):
+        self._body = body_bytes
+        self._headers = headers
+        self._cache = {}
+        self.status = 200
+        self.reason = "OK"
+        self._url = url
+
+
+class MockTransport(AsyncHttpTransport):
+    async def send(self, request: HttpRequest, **kwargs) -> AioHttpTransportResponse:
+        if request.url.endswith("notfound.png"):
+            raise ResourceNotFoundError(MockAiohttpClientResponse404(request.url, b""))
+        else:
+            return AioHttpTransportResponse(
+                request,
+                MockAiohttpClientResponse(
+                    request.url,
+                    b"test content",
+                    {
+                        "Content-Type": "application/octet-stream",
+                        "Content-Range": "bytes 0-27/28",
+                        "Content-Length": "28",
+                    },
+                ),
+            )
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        pass
+
+    async def open(self):
+        pass
+
+    async def close(self):
+        pass
+
+
 class MockAsyncPageIterator:
     def __init__(self, data):
         self.data = data

tests/snapshots/test_app/test_ask_vision/client0/result.json

@@ -8,7 +8,7 @@
                 "Financial Market Analysis Report 2023.pdf#page=2"
             ],
             "images": [
-                "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z/C/HgAGgwJ/lK3Q6wAAAABJRU5ErkJggg=="
+                "data:image/png;base64,dGVzdCBjb250ZW50"
             ],
             "text": [
                 "Financial Market Analysis Report 2023.pdf#page=7:  This section examines the correlations between stock indices, cryptocurrency prices, and commodity prices, revealing how changes in one market can have ripple effects across the financial ecosystem.### Impact of Macroeconomic Factors   <figure><figcaption>Impact of Interest Rates, Inflation, and GDP Growth on Financial Markets<br>The image is a line graph titled \"on Financial Markets\" displaying data from 2018 to 2023. It tracks three variables: Interest Rates %, Inflation Data %, and GDP Growth %, each represented by a different colored line (blue for Interest Rates, orange for Inflation Data, and gray for GDP Growth). Interest Rates % start around 2% in 2018, dip to about 0.25% in 2021, then rise to 1.5% in 2023. Inflation Data % begin at approximately 1.9% in 2018, rise to a peak near 3.4% in 2022, and then decrease to 2.5% in 2023. GDP Growth % shows significant fluctuations, starting at 3% in 2018, plunging to almost -4% in 2020, then rebounding to around 4.5% in 2021 before gradually declining to around 2.8% in 2023.</figcaption></figure>   Macroeconomic factors such as interest rates, inflation, and GDP growth play a pivotal role in shaping financial markets.",
@@ -124,7 +124,7 @@
                             },
                             {
                                 "image_url": {
-                                    "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z/C/HgAGgwJ/lK3Q6wAAAABJRU5ErkJggg=="
+                                    "url": "data:image/png;base64,dGVzdCBjb250ZW50"
                                 },
                                 "type": "image_url"
                             },