From 934b1d2c0dfe5eb62b23ac1b944d070848133f47 Mon Sep 17 00:00:00 2001 From: Josh Bradley Date: Fri, 27 Dec 2024 20:54:39 -0500 Subject: [PATCH] add more unit tests --- backend/src/api/common.py | 10 +++--- backend/tests/conftest.py | 10 +++--- backend/tests/unit/test_common.py | 51 ++++++++++++++++++++++--------- 3 files changed, 47 insertions(+), 24 deletions(-) diff --git a/backend/src/api/common.py b/backend/src/api/common.py index 1b81ba5d..bec50933 100644 --- a/backend/src/api/common.py +++ b/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 diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py index b810e2a5..cff312bf 100644 --- a/backend/tests/conftest.py +++ b/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 diff --git a/backend/tests/unit/test_common.py b/backend/tests/unit/test_common.py index a86d06ee..4b6b078e 100644 --- a/backend/tests/unit/test_common.py +++ b/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" + )