add more type hints

Author: jgbradley1

backend/tests/conftest.py

@@ -15,7 +15,7 @@
 
 
 @pytest.fixture(scope="session")
-def blob_with_data_container_name(blob_service_client):
+def blob_with_data_container_name(blob_service_client: BlobServiceClient):
     # create a storage container and upload some data
     container_name = "container-with-data"
     blob_service_client.create_container(container_name)
@@ -27,7 +27,7 @@ def blob_with_data_container_name(blob_service_client):
 
 
 @pytest.fixture(scope="session")
-def blob_service_client():
+def blob_service_client() -> Generator[BlobServiceClient, None, None]:
     blob_service_client = BlobServiceClient.from_connection_string(
         os.environ["STORAGE_CONNECTION_STRING"]
     )
@@ -36,7 +36,7 @@ def blob_service_client():
 
 
 @pytest.fixture(scope="session")
-def cosmos_client():
+def cosmos_client() -> Generator[CosmosClient, None, None]:
     """ "Initializes the CosmosDB databases that graphrag expects at startup time."""
     # setup
     client = CosmosClient.from_connection_string(os.environ["COSMOS_CONNECTION_STRING"])
@@ -53,7 +53,9 @@ def cosmos_client():
 
 
 @pytest.fixture(scope="session")
-def container_with_graphml_file(blob_service_client, cosmos_client):
+def container_with_graphml_file(
+    blob_service_client: BlobServiceClient, cosmos_client: CosmosClient
+):
     """create a storage container and upload a fake graphml file"""
     container_name = "container-with-graphml"
     sanitized_name = sanitize_name(container_name)
@@ -79,7 +81,9 @@ def container_with_graphml_file(blob_service_client, cosmos_client):
 
 
 @pytest.fixture(scope="session")
-def container_with_index_files(blob_service_client, cosmos_client):
+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"""
     container_name = "container-with-index-files"
     sanitized_name = sanitize_name(container_name)
@@ -122,6 +126,6 @@ def container_with_index_files(blob_service_client, cosmos_client):
 
 
 @pytest.fixture(scope="session")
-def client(request) -> Generator:
+def client(request) -> Generator[TestClient, None, None]:
     with TestClient(app) as c:
         yield c

backend/tests/integration/test_api_data.py

@@ -6,8 +6,10 @@
 
 import os
 
+from azure.cosmos import CosmosClient
 
-def test_upload_files(cosmos_client, client):
+
+def test_upload_files(cosmos_client: CosmosClient, client):
     """Test uploading files to a data blob container."""
     # create a single file
     with open("test.txt", "wb") as f:
@@ -26,14 +28,14 @@ def test_upload_files(cosmos_client, client):
         os.remove("test.txt")
 
 
-def test_delete_files(cosmos_client, client):
+def test_delete_files(cosmos_client: CosmosClient, client):
     """Test deleting a data blob container."""
     # delete a data blob container
     response = client.delete("/data/testContainer")
     assert response.status_code == 200
 
 
-def test_get_list_of_data_containers(cosmos_client, client):
+def test_get_list_of_data_containers(cosmos_client: CosmosClient, client):
     """Test getting a list of all data blob containers."""
     response = client.get("/data")
     assert response.status_code == 200

backend/tests/integration/test_api_graph.py

@@ -5,7 +5,7 @@
 """
 
 
-def test_get_graphml_file(client, container_with_graphml_file):
+def test_get_graphml_file(client, container_with_graphml_file: str):
     """Test retrieving a graphml file endpoint."""
     url = f"/graph/graphml/{container_with_graphml_file}"
     response = client.get(url)

backend/tests/integration/test_api_index.py

@@ -4,14 +4,16 @@
 Integration tests for the /index API endpoints.
 """
 
+from azure.cosmos import CosmosClient
 
-def test_get_list_of_index_containers_empty(client, cosmos_client):
+
+def test_get_list_of_index_containers_empty(client, cosmos_client: CosmosClient):
     """Test getting a list of all blob containers holding an index."""
     response = client.get("/index")
     assert response.status_code == 200
 
 
-def test_schedule_index_without_data(client, cosmos_client):
+def test_schedule_index_without_data(client, cosmos_client: CosmosClient):
     """Test scheduling an index job with a non-existent data blob container."""
     response = client.post(
         "/index",