Add more tests for blobmanger

Author: pamelafox

app/backend/app.py

@@ -357,13 +357,17 @@ async def upload(auth_claims: dict[str, Any]):
     if "file" not in request_files:
         return jsonify({"message": "No file part in the request", "status": "failed"}), 400
 
-    user_oid = auth_claims["oid"]
-    file = request_files.getlist("file")[0]
-    adls_manager: AdlsBlobManager = current_app.config[CONFIG_USER_BLOB_MANAGER]
-    file_url = await adls_manager.upload_blob(file, file.filename, user_oid)
-    ingester: UploadUserFileStrategy = current_app.config[CONFIG_INGESTER]
-    await ingester.add_file(File(content=file, url=file_url, acls={"oids": [user_oid]}), user_oid=user_oid)
-    return jsonify({"message": "File uploaded successfully"}), 200
+    try:
+        user_oid = auth_claims["oid"]
+        file = request_files.getlist("file")[0]
+        adls_manager: AdlsBlobManager = current_app.config[CONFIG_USER_BLOB_MANAGER]
+        file_url = await adls_manager.upload_blob(file, file.filename, user_oid)
+        ingester: UploadUserFileStrategy = current_app.config[CONFIG_INGESTER]
+        await ingester.add_file(File(content=file, url=file_url, acls={"oids": [user_oid]}), user_oid=user_oid)
+        return jsonify({"message": "File uploaded successfully"}), 200
+    except Exception as error:
+        current_app.logger.error("Error uploading file: %s", error)
+        return jsonify({"message": "Error uploading file, check server logs for details.", "status": "failed"}), 500
 
 
 @bp.post("/delete_uploaded")

app/backend/prepdocslib/blobmanager.py

@@ -38,10 +38,6 @@ def sourcepage_from_file_page(cls, filename, page=0) -> str:
         else:
             return os.path.basename(filename)
 
-    @classmethod
-    def blob_image_name_from_file_page(cls, filename, page=0) -> str:
-        return os.path.splitext(os.path.basename(filename))[0] + f"-{page+1}" + ".png"
-
     @classmethod
     def blob_name_from_file_name(cls, filename) -> str:
         return os.path.basename(filename)

tests/conftest.py

@@ -308,6 +308,14 @@ async def mock_exists(*args, **kwargs):
     monkeypatch.setattr("azure.storage.blob.aio.ContainerClient.exists", mock_exists)
 
 
+@pytest.fixture
+def mock_blob_container_client_does_not_exist(monkeypatch):
+    async def mock_exists(*args, **kwargs):
+        return False
+
+    monkeypatch.setattr("azure.storage.blob.aio.ContainerClient.exists", mock_exists)
+
+
 envs = [
     {
         "OPENAI_HOST": "openai",

tests/test_blob_manager.py

@@ -25,17 +25,12 @@ def blob_manager(monkeypatch):
 
 @pytest.mark.asyncio
 @pytest.mark.skipif(sys.version_info.minor < 10, reason="requires Python 3.10 or higher")
-async def test_upload_and_remove(monkeypatch, mock_env, blob_manager):
+async def test_upload_and_remove(monkeypatch, mock_env, mock_blob_container_client_exists, blob_manager):
     with NamedTemporaryFile(suffix=".pdf") as temp_file:
         f = File(temp_file.file)
         filename = os.path.basename(f.content.name)
 
-        # Set up mocks used by upload_blob
-        async def mock_exists(*args, **kwargs):
-            return True
-
-        monkeypatch.setattr("azure.storage.blob.aio.ContainerClient.exists", mock_exists)
-
+        # Set up mock of upload_blob
         async def mock_upload_blob(self, name, *args, **kwargs):
             assert name == filename
             return azure.storage.blob.aio.BlobClient.from_blob_url(
@@ -78,17 +73,12 @@ async def mock_delete_blob(self, name, *args, **kwargs):
 
 @pytest.mark.asyncio
 @pytest.mark.skipif(sys.version_info.minor < 10, reason="requires Python 3.10 or higher")
-async def test_upload_and_remove_all(monkeypatch, mock_env, blob_manager):
+async def test_upload_and_remove_all(monkeypatch, mock_env, mock_blob_container_client_exists, blob_manager):
     with NamedTemporaryFile(suffix=".pdf") as temp_file:
         f = File(temp_file.file)
         filename = os.path.basename(f.content.name)
 
-        # Set up mocks used by upload_blob
-        async def mock_exists(*args, **kwargs):
-            return True
-
-        monkeypatch.setattr("azure.storage.blob.aio.ContainerClient.exists", mock_exists)
-
+        # Set up mock of upload_blob
         async def mock_upload_blob(self, name, *args, **kwargs):
             assert name == filename
             return azure.storage.blob.aio.BlobClient.from_blob_url(
@@ -161,12 +151,9 @@ async def mock_upload_blob(self, name, *args, **kwargs):
 
 @pytest.mark.asyncio
 @pytest.mark.skipif(sys.version_info.minor < 10, reason="requires Python 3.10 or higher")
-async def test_dont_remove_if_no_container(monkeypatch, mock_env, blob_manager):
-    async def mock_exists(*args, **kwargs):
-        return False
-
-    monkeypatch.setattr("azure.storage.blob.aio.ContainerClient.exists", mock_exists)
-
+async def test_dont_remove_if_no_container(
+    monkeypatch, mock_env, mock_blob_container_client_does_not_exist, blob_manager
+):
     async def mock_delete_blob(*args, **kwargs):
         assert False, "delete_blob() shouldn't have been called"
 
@@ -177,7 +164,8 @@ async def mock_delete_blob(*args, **kwargs):
 
 @pytest.mark.asyncio
 @pytest.mark.skipif(sys.version_info.minor < 10, reason="requires Python 3.10 or higher")
-async def test_upload_document_image(monkeypatch, mock_env):
+@pytest.mark.parametrize("directory_exists", [True, False])
+async def test_upload_document_image(monkeypatch, mock_env, directory_exists):
     # Create a blob manager with an image container
     blob_manager = BlobManager(
         endpoint=f"https://{os.environ['AZURE_STORAGE_ACCOUNT']}.blob.core.windows.net",
@@ -202,10 +190,15 @@ async def test_upload_document_image(monkeypatch, mock_env):
 
         # Mock container client operations
         async def mock_exists(*args, **kwargs):
-            return True
+            return directory_exists
 
         monkeypatch.setattr("azure.storage.blob.aio.ContainerClient.exists", mock_exists)
 
+        async def mock_create_container(*args, **kwargs):
+            return
+
+        monkeypatch.setattr("azure.storage.blob.aio.ContainerClient.create_container", mock_create_container)
+
         expected_blob_name = f"{os.path.basename(temp_file.name)}/page{image_page_num}/{image_filename}"
 
         async def mock_upload_blob(self, name, *args, **kwargs):
@@ -241,3 +234,80 @@ def test_sourcepage_from_file_page():
 def test_blob_name_from_file_name():
     assert BlobManager.blob_name_from_file_name("tmp/test.pdf") == "test.pdf"
     assert BlobManager.blob_name_from_file_name("tmp/test.html") == "test.html"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(sys.version_info.minor < 10, reason="requires Python 3.10 or higher")
+async def test_download_blob(monkeypatch, mock_env, mock_blob_container_client_exists, blob_manager):
+    # Mock the download_blob method
+    test_content = b"test content bytes"
+
+    class MockDownloadResponse:
+        def __init__(self):
+            # Create properties with content_settings
+            class ContentSettings:
+                content_type = "application/pdf"
+
+            class Properties:
+                def __init__(self):
+                    self.content_settings = ContentSettings()
+
+            self.properties = Properties()
+
+        async def readall(self):
+            return test_content
+
+    async def mock_download_blob(*args, **kwargs):
+        return MockDownloadResponse()
+
+    monkeypatch.setattr("azure.storage.blob.aio.BlobClient.download_blob", mock_download_blob)
+
+    result = await blob_manager.download_blob("test_document.pdf")
+
+    assert result is not None
+    content, properties = result
+    assert content == test_content
+    assert properties["content_settings"]["content_type"] == "application/pdf"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(sys.version_info.minor < 10, reason="requires Python 3.10 or higher")
+async def test_download_blob_not_found(monkeypatch, mock_env, mock_blob_container_client_exists, blob_manager):
+    # Mock the download_blob method to raise ResourceNotFoundError
+    async def mock_download_blob(*args, **kwargs):
+        from azure.core.exceptions import ResourceNotFoundError
+
+        raise ResourceNotFoundError("Blob not found")
+
+    monkeypatch.setattr("azure.storage.blob.aio.BlobClient.download_blob", mock_download_blob)
+
+    result = await blob_manager.download_blob("nonexistent.pdf")
+
+    assert result is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(sys.version_info.minor < 10, reason="requires Python 3.10 or higher")
+async def test_download_blob_container_not_exist(
+    monkeypatch, mock_env, mock_blob_container_client_does_not_exist, blob_manager
+):
+    result = await blob_manager.download_blob("test_document.pdf")
+
+    assert result is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(sys.version_info.minor < 10, reason="requires Python 3.10 or higher")
+async def test_download_blob_empty_path(monkeypatch, mock_env, mock_blob_container_client_exists, blob_manager):
+    result = await blob_manager.download_blob("")
+
+    assert result is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(sys.version_info.minor < 10, reason="requires Python 3.10 or higher")
+async def test_download_blob_with_user_oid(monkeypatch, mock_env, blob_manager):
+    with pytest.raises(ValueError) as excinfo:
+        await blob_manager.download_blob("test_document.pdf", user_oid="user123")
+
+    assert "user_oid is not supported for BlobManager" in str(excinfo.value)

tests/test_upload.py

@@ -18,7 +18,6 @@
 from .mocks import MockClient, MockEmbeddingsClient
 
 
-# parameterize for directory existing or not
 @pytest.mark.asyncio
 @pytest.mark.parametrize("directory_exists", [True, False])
 async def test_upload_file(auth_client, monkeypatch, mock_data_lake_service_client, directory_exists):
@@ -115,6 +114,38 @@ async def mock_upload_documents(self, documents):
     assert directory_created[0] == (not directory_exists)
 
 
+@pytest.mark.asyncio
+async def test_upload_file_error_wrong_directory_owner(auth_client, monkeypatch, mock_data_lake_service_client):
+
+    # Create a mock class for DataLakeDirectoryClient that includes the _client attribute
+    class MockDataLakeDirectoryClient:
+        def __init__(self, *args, **kwargs):
+            self._client = object()
+            self.url = "https://test.blob.core.windows.net/container/path"
+
+        async def get_directory_properties(self, *args, **kwargs):
+            return {"name": "test-directory"}
+
+        async def get_access_control(self, *args, **kwargs):
+            return {"owner": "OID_Y"}
+
+    # Replace the DataLakeDirectoryClient with our mock
+    monkeypatch.setattr(
+        azure.storage.filedatalake.aio.FileSystemClient,
+        "get_directory_client",
+        lambda *args, **kwargs: MockDataLakeDirectoryClient(),
+    )
+
+    response = await auth_client.post(
+        "/upload",
+        headers={"Authorization": "Bearer test"},
+        files={"file": FileStorage(BytesIO(b"foo;bar"), filename="a.txt")},
+    )
+    message = (await response.get_json())["message"]
+    assert message == "Error uploading file, check server logs for details."
+    assert response.status_code == 500
+
+
 @pytest.mark.asyncio
 async def test_list_uploaded(auth_client, monkeypatch, mock_data_lake_service_client):
     response = await auth_client.get("/list_uploaded", headers={"Authorization": "Bearer test"})

todo.txt

@@ -4,7 +4,6 @@ TODO:
    * DocIntelligence does not support JSON/MD - note in docs
    * Verbalization strategy from Sri's skillset
    * Filter images to match page numbers, using shaping + conditional skills
-* Update all TODOs in the code/docs
 * Fix/add unit tests - check coverage