add more unit tests

backend/src/api/common.py

@@ -176,11 +176,13 @@ def retrieve_original_blob_container_name(sanitized_name: str) -> str | None:
         container_store_client = azure_client_manager.get_cosmos_container_client(
             database="graphrag", container="container-store"
         )
-        for item in container_store_client.read_all_items():
-            if item["id"] == sanitized_name:
-                return item["human_readable_name"]
+        try:
+            return container_store_client.read_item(sanitized_name, sanitized_name)[
+                "human_readable_name"
+            ]
+        except exceptions.CosmosResourceNotFoundError:
+            return None
     except Exception:
         raise HTTPException(
             status_code=500, detail="Error retrieving original blob name."
         )
-    return None

backend/tests/conftest.py

@@ -37,7 +37,7 @@ def blob_service_client() -> Generator[BlobServiceClient, None, None]:
 
 @pytest.fixture(scope="session")
 def cosmos_client() -> Generator[CosmosClient, None, None]:
-    """ "Initializes the CosmosDB databases that graphrag expects at startup time."""
+    """Initializes the CosmosDB databases that graphrag expects at startup time."""
     # setup
     client = CosmosClient.from_connection_string(os.environ["COSMOS_CONNECTION_STRING"])
     db_client = client.create_database_if_not_exists(id="graphrag")
@@ -56,7 +56,7 @@ def cosmos_client() -> Generator[CosmosClient, None, None]:
 def container_with_graphml_file(
     blob_service_client: BlobServiceClient, cosmos_client: CosmosClient
 ):
-    """create a storage container and upload a fake graphml file"""
+    """Create a storage container that mimics a valid index and upload a fake graphml file"""
     container_name = "container-with-graphml"
     sanitized_name = sanitize_name(container_name)
     if not blob_service_client.get_container_client(sanitized_name).exists():
@@ -71,7 +71,7 @@ def container_with_graphml_file(
     ).get_container_client("container-store")
     container_store_client.upsert_item({
         "id": sanitized_name,
-        "human_readable_name": sanitized_name,
+        "human_readable_name": container_name,
         "type": "index",
     })
     yield container_name
@@ -84,7 +84,7 @@ def container_with_graphml_file(
 def container_with_index_files(
     blob_service_client: BlobServiceClient, cosmos_client: CosmosClient
 ):
-    """create a storage container and upload a set of parquet files associated with a valid index"""
+    """Create a storage container and upload a set of parquet files associated with a valid index"""
     container_name = "container-with-index-files"
     sanitized_name = sanitize_name(container_name)
     if not blob_service_client.get_container_client(sanitized_name).exists():
@@ -116,7 +116,7 @@ def container_with_index_files(
     ).get_container_client("container-store")
     container_store_client.upsert_item({
         "id": sanitized_name,
-        "human_readable_name": sanitized_name,
+        "human_readable_name": container_name,
         "type": "index",
     })
     yield container_name

backend/tests/unit/test_common.py

@@ -4,7 +4,10 @@
 import pytest
 
 from src.api.common import (
+    retrieve_original_blob_container_name,
+    sanitize_name,
     validate_blob_container_name,
+    validate_index_file_exist,
 )
 
 
@@ -15,26 +18,44 @@ def test_validate_blob_container_name():
     # test invalid container name
     with pytest.raises(ValueError):
         validate_blob_container_name("invalidContainerName")
+    with pytest.raises(ValueError):
+        validate_blob_container_name(
+            "invalidcontainernameinvalidcontainernameinvalidcontainerinvalids"
+        )
     with pytest.raises(ValueError):
         validate_blob_container_name("*invalidContainerName")
     with pytest.raises(ValueError):
         validate_blob_container_name("invalid+ContainerName")
     with pytest.raises(ValueError):
-        validate_blob_container_name("invalid--ContainerName")
+        validate_blob_container_name("invalid--containername")
     with pytest.raises(ValueError):
-        validate_blob_container_name("invalidContainerName-")
+        validate_blob_container_name("invalidcontainername-")
+
+
+def test_retrieve_original_blob_container_name(container_with_graphml_file):
+    """Test the src.api.common.retrieve_original_blob_container_name function."""
+    # test retrieving a valid container name
+    original_name = container_with_graphml_file
+    sanitized_name = sanitize_name(original_name)
+    assert retrieve_original_blob_container_name(sanitized_name) == original_name
+    # test retrieving an invalid container name
+    assert retrieve_original_blob_container_name("nonexistent-container") is None
 
 
-# def test_validate_index_file_exist():
-#     """Test the src.api.common.validate_index_file_exist function."""
-#     # test valid index and file
-#     assert validate_index_file_exist("validindex", "validfile") is None
-#     # test invalid index
-#     with pytest.raises(ValueError):
-#         validate_index_file_exist("invalidindex", "validfile")
-#     # test invalid file
-#     with pytest.raises(ValueError):
-#         validate_index_file_exist("validindex", "invalidfile")
-#     # test invalid index and file
-#     with pytest.raises(ValueError):
-#         validate_index_file_exist("invalidindex", "invalidfile")
+def test_validate_index_file_exist(container_with_graphml_file):
+    """Test the src.api.common.validate_index_file_exist function."""
+    original_name = container_with_graphml_file
+    sanitized_name = sanitize_name(original_name)
+    # test with a valid index and valid file
+    assert (
+        validate_index_file_exist(sanitized_name, "output/summarized_graph.graphml")
+        is None
+    )
+    # test with a valid index and non-existent file
+    with pytest.raises(ValueError):
+        validate_index_file_exist(sanitized_name, "non-existent-file")
+    # test non-existent index and valid file
+    with pytest.raises(ValueError):
+        validate_index_file_exist(
+            "nonexistent-index", "output/summarized_graph.graphml"
+        )