Fix Python CI test failures and improve code quality

Co-authored-by: jgbradley1 <654554+jgbradley1@users.noreply.github.com>

Author: Copilot

graphrag/storage/factory.py

@@ -5,8 +5,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
-from typing import TYPE_CHECKING, Any, ClassVar
+from contextlib import suppress
+from typing import TYPE_CHECKING, ClassVar
 
 from graphrag.config.enums import OutputType
 from graphrag.storage.blob_pipeline_storage import create_blob_storage
@@ -15,6 +15,8 @@
 from graphrag.storage.memory_pipeline_storage import MemoryPipelineStorage
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from graphrag.storage.pipeline_storage import PipelineStorage
 
 
@@ -31,47 +33,52 @@ class StorageFactory:
     storage_types: ClassVar[dict[str, type]] = {}  # For backward compatibility
 
     @classmethod
-    def register(cls, storage_type: str, creator: Callable[..., PipelineStorage]) -> None:
+    def register(
+        cls, storage_type: str, creator: Callable[..., PipelineStorage]
+    ) -> None:
         """Register a custom storage implementation.
-        
+
         Args:
             storage_type: The type identifier for the storage.
             creator: A callable that creates an instance of the storage.
         """
         cls._storage_registry[storage_type] = creator
-        
+
         # For backward compatibility with code that may access storage_types directly
-        if callable(creator) and hasattr(creator, "__annotations__") and "return" in creator.__annotations__:
-            try:
+        if (
+            callable(creator)
+            and hasattr(creator, "__annotations__")
+            and "return" in creator.__annotations__
+        ):
+            with suppress(TypeError, KeyError):
                 cls.storage_types[storage_type] = creator.__annotations__["return"]
-            except (TypeError, KeyError):
-                # Just ignore if we can't maintain backward compatibility in this case
-                pass
 
     @classmethod
     def create_storage(
         cls, storage_type: OutputType | str, kwargs: dict
     ) -> PipelineStorage:
         """Create a storage object from the provided type.
-        
+
         Args:
             storage_type: The type of storage to create.
             kwargs: Additional keyword arguments for the storage constructor.
-            
+
         Returns
         -------
             A PipelineStorage instance.
-        
+
         Raises
         ------
             ValueError: If the storage type is not registered.
         """
-        storage_type_str = storage_type.value if isinstance(storage_type, OutputType) else storage_type
-        
+        storage_type_str = (
+            storage_type.value if isinstance(storage_type, OutputType) else storage_type
+        )
+
         if storage_type_str not in cls._storage_registry:
             msg = f"Unknown storage type: {storage_type}"
             raise ValueError(msg)
-            
+
         return cls._storage_registry[storage_type_str](**kwargs)
 
     @classmethod
@@ -89,4 +96,4 @@ def is_supported_storage_type(cls, storage_type: str) -> bool:
 StorageFactory.register(OutputType.blob.value, create_blob_storage)
 StorageFactory.register(OutputType.cosmosdb.value, create_cosmosdb_storage)
 StorageFactory.register(OutputType.file.value, create_file_storage)
-StorageFactory.register(OutputType.memory.value, lambda **kwargs: MemoryPipelineStorage())
+StorageFactory.register(OutputType.memory.value, lambda **_: MemoryPipelineStorage())

tests/integration/storage/test_factory.py

@@ -15,6 +15,7 @@
 from graphrag.storage.factory import StorageFactory
 from graphrag.storage.file_pipeline_storage import FilePipelineStorage
 from graphrag.storage.memory_pipeline_storage import MemoryPipelineStorage
+from graphrag.storage.pipeline_storage import PipelineStorage
 
 # cspell:disable-next-line well-known-key
 WELL_KNOWN_BLOB_STORAGE_KEY = "DefaultEndpointsProtocol=http;AccountName=devstoreaccount1;AccountKey=Eby8vdM02xNOcqFlqUwJPLlmEtlCDXJ1OUzFT50uSRZ6IFsuFq2UVErCz4I6tq/K1SZFPTOtr/KBHBeksoGMGw==;BlobEndpoint=http://127.0.0.1:10000/devstoreaccount1;"
@@ -62,15 +63,25 @@ def test_create_memory_storage():
 
 
 def test_register_and_create_custom_storage():
-    class CustomStorage:
-        def __init__(self, **kwargs):
-            self.initialized = True
-
-    StorageFactory.register("custom", lambda **kwargs: CustomStorage(**kwargs))
+    """Test registering and creating a custom storage type."""
+    from unittest.mock import MagicMock
+
+    # Create a mock that satisfies the PipelineStorage interface
+    custom_storage_class = MagicMock(spec=PipelineStorage)
+    # Make the mock return a mock instance when instantiated
+    instance = MagicMock()
+    # We can set attributes on the mock instance, even if they don't exist on PipelineStorage
+    instance.initialized = True
+    custom_storage_class.return_value = instance
+
+    StorageFactory.register("custom", lambda **kwargs: custom_storage_class(**kwargs))
     storage = StorageFactory.create_storage("custom", {})
-    assert isinstance(storage, CustomStorage)
-    assert storage.initialized
-    
+
+    assert custom_storage_class.called
+    assert storage is instance
+    # Access the attribute we set on our mock
+    assert storage.initialized is True  # type: ignore # Attribute only exists on our mock
+
     # Check if it's in the list of registered storage types
     assert "custom" in StorageFactory.get_storage_types()
     assert StorageFactory.is_supported_storage_type("custom")