update pytest to use bacj engine

Author: jgbradley1

graphrag/storage/blob_pipeline_storage.py

@@ -65,11 +65,11 @@ def __init__(
             self._container_name,
             self._path_prefix,
         )
-        self.create_container()
+        self._create_container()
 
-    def create_container(self) -> None:
+    def _create_container(self) -> None:
         """Create the container if it does not exist."""
-        if not self.container_exists():
+        if not self._container_exists():
             container_name = self._container_name
             container_names = [
                 container.name
@@ -78,12 +78,12 @@ def create_container(self) -> None:
             if container_name not in container_names:
                 self._blob_service_client.create_container(container_name)
 
-    def delete_container(self) -> None:
+    def _delete_container(self) -> None:
         """Delete the container."""
-        if self.container_exists():
+        if self._container_exists():
             self._blob_service_client.delete_container(self._container_name)
 
-    def container_exists(self) -> bool:
+    def _container_exists(self) -> bool:
         """Check if the container exists."""
         container_name = self._container_name
         container_names = [
@@ -119,7 +119,7 @@ def find(
             file_pattern.pattern,
         )
 
-        def blobname(blob_name: str) -> str:
+        def _blobname(blob_name: str) -> str:
             if blob_name.startswith(self._path_prefix):
                 blob_name = blob_name.replace(self._path_prefix, "", 1)
             if blob_name.startswith("/"):
@@ -146,7 +146,7 @@ def item_filter(item: dict[str, Any]) -> bool:
                 if match and blob.name.startswith(base_dir):
                     group = match.groupdict()
                     if item_filter(group):
-                        yield (blobname(blob.name), group)
+                        yield (_blobname(blob.name), group)
                         num_loaded += 1
                         if max_count > 0 and num_loaded >= max_count:
                             break
@@ -203,7 +203,7 @@ async def set(self, key: str, value: Any, encoding: str | None = None) -> None:
         except Exception:
             log.exception("Error setting key %s: %s", key)
 
-    def set_df_json(self, key: str, dataframe: Any) -> None:
+    def _set_df_json(self, key: str, dataframe: Any) -> None:
         """Set a json dataframe."""
         if self._connection_string is None and self._storage_account_name:
             dataframe.to_json(
@@ -225,7 +225,7 @@ def set_df_json(self, key: str, dataframe: Any) -> None:
                 force_ascii=False,
             )
 
-    def set_df_parquet(self, key: str, dataframe: Any) -> None:
+    def _set_df_parquet(self, key: str, dataframe: Any) -> None:
         """Set a parquet dataframe."""
         if self._connection_string is None and self._storage_account_name:
             dataframe.to_parquet(

graphrag/storage/file_pipeline_storage.py

@@ -47,7 +47,6 @@ def find(
         def item_filter(item: dict[str, Any]) -> bool:
             if file_filter is None:
                 return True
-
             return all(re.match(value, item[key]) for key, value in file_filter.items())
 
         search_path = Path(self._root_dir) / (base_dir or "")

tests/integration/storage/test_blob_pipeline_storage.py

@@ -44,7 +44,7 @@ async def test_find():
             output = await storage.get("test.txt")
             assert output is None
     finally:
-        storage.delete_container()
+        storage._delete_container()  # noqa: SLF001
 
 
 async def test_dotprefix():
@@ -59,7 +59,7 @@ async def test_dotprefix():
         items = [item[0] for item in items]
         assert items == ["input/christmas.txt"]
     finally:
-        storage.delete_container()
+        storage._delete_container()  # noqa: SLF001
 
 
 async def test_child():
@@ -93,4 +93,4 @@ async def test_child():
             has_test = await parent.has("input/test.txt")
             assert not has_test
     finally:
-        parent.delete_container()
+        parent._delete_container()  # noqa: SLF001

tests/smoke/test_fixtures.py

@@ -100,8 +100,8 @@ async def prepare_azurite_data(input_path: str, azure: dict) -> Callable[[], Non
         container_name=input_container,
     )
     # Bounce the container if it exists to clear out old run data
-    input_storage.delete_container()
-    input_storage.create_container()
+    input_storage._delete_container()  # noqa: SLF001
+    input_storage._create_container()  # noqa: SLF001
 
     # Upload data files
     txt_files = list((root / "input").glob("*.txt"))
@@ -116,7 +116,7 @@ async def prepare_azurite_data(input_path: str, azure: dict) -> Callable[[], Non
         )
         await input_storage.set(file_path, text, encoding="utf-8")
 
-    return lambda: input_storage.delete_container()
+    return lambda: input_storage._delete_container()  # noqa: SLF001
 
 
 class TestIndexer: