Added tests for memory_pipeline_storage

Author: Snehit Gajjar

graphrag/storage/memory_pipeline_storage.py

@@ -3,7 +3,6 @@
 
 """A module containing 'InMemoryStorage' model."""
 
-import asyncio
 import logging
 import re
 from collections.abc import Iterator
@@ -36,7 +35,8 @@ def find(
         """Find files in the storage using a file pattern, as well as a custom filter function."""
         # In-memory storage does not support file finding
         results: list[tuple[str, dict[str, Any]]] = []
-        results = [(key, {}) for key in self._storage["files"]]
+        if self._storage.get("files") is not None:
+            results = [(key, {}) for key in self._storage["files"]]
         return iter(results)
 
     async def get(
@@ -84,7 +84,7 @@ async def delete(self, key: str) -> None:
         Args:
             - key - The key to delete.
         """
-        del self._storage[key]
+        del self._storage["files"][key]
 
     async def clear(self) -> None:
         """Clear the storage."""
@@ -109,14 +109,19 @@ def set_sync(self, key: str, value: Any, encoding: str | None = None) -> None:
             - key - The key to set the value for.
             - value - The value to set.
         """
-        task = asyncio.create_task(self.set(key, value, encoding))
-        logger.info("Setting value for key '%s' in memory storage", task.get_name())
+        logger.info("Setting value for key '%s' in memory storage", key)
+        if "files" not in self._storage:
+            self._storage["files"] = {}
+        self._storage["files"][key] = value
 
 
 def create_memory_storage(**kwargs: Any) -> PipelineStorage:
     """Create a memory based storage."""
     logger.info("Creating memory storage")
     memorystorage = MemoryPipelineStorage()
+    if kwargs.get("input_files") is None:
+        return memorystorage
+
     for key, value in kwargs["input_files"].items():
         memorystorage.set_sync(key, value)
     return memorystorage

pyproject.toml

@@ -109,6 +109,7 @@ mkdocs-jupyter = "^0.25.1"
 mkdocs-exclude-search = "^0.6.6"
 pytest-dotenv = "^0.5.2"
 mkdocs-typer = "^0.0.3"
+freezegun = "^1.5.3"
 
 [build-system]
 requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"]
@@ -268,3 +269,6 @@ asyncio_default_fixture_loop_scope = "function"
 asyncio_mode = "auto"
 timeout = 1000
 env_files = [".env"]
+pythonpath = [
+  "./tests"
+]

tests/integration/storage/test_memory_pipeline_storage.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+"""Blob Storage Tests."""
+
+import os
+import re
+from datetime import datetime, timezone
+
+import pandas as pd
+from freezegun import freeze_time
+
+from graphrag.storage.memory_pipeline_storage import (
+    MemoryPipelineStorage,
+    create_memory_storage,
+)
+
+__dirname__ = os.path.dirname(__file__)
+
+
+async def test_find():
+    input_files = {
+        "input_files": {
+            "tests/fixtures/text/input/dulce.txt": pd.DataFrame({
+                "text": ["Dulce et decorum est"],
+                "creation_date": [datetime(2023, 1, 1, tzinfo=timezone.utc)],
+            }),
+        }
+    }
+
+    storage = create_memory_storage(**input_files)
+
+    items = list(
+        storage.find(
+            base_dir="any/path",
+            file_pattern=re.compile(r".*\.txt$"),
+            file_filter=None,
+        )
+    )
+    assert items == [("tests/fixtures/text/input/dulce.txt", {})]
+    output = await storage.get("tests/fixtures/text/input/dulce.txt")
+    assert len(output) > 0
+
+    await storage.set("test.txt", "Hello, World!", encoding="utf-8")
+    output = await storage.get("test.txt")
+    assert output == "Hello, World!"
+    await storage.delete("test.txt")
+    output = await storage.get("test.txt")
+    assert output is None
+
+
+@freeze_time("2023-10-01T15:55:00.12345")
+async def test_get_creation_date():
+    storage = MemoryPipelineStorage()
+
+    creation_date = await storage.get_creation_date(
+        "tests/fixtures/text/input/dulce.txt"
+    )
+
+    parsed_datetime = datetime.now(timezone.utc).isoformat()
+
+    assert parsed_datetime == creation_date
+
+
+def test_child():
+    input_files = {
+        "input_files": {
+            "tests/fixtures/text/input/dulce.txt": pd.DataFrame({
+                "text": ["Dulce et decorum est"],
+                "creation_date": [datetime(2023, 1, 1, tzinfo=timezone.utc)],
+            }),
+        }
+    }
+
+    storage = create_memory_storage(**input_files)
+
+    child_storage = storage.child("tests/fixtures/text/input")
+    items = list(
+        storage.find(base_dir="any/path", file_pattern=re.compile(r".*\.txt$"))
+    )
+    assert items == [("tests/fixtures/text/input/dulce.txt", {})]
+    assert child_storage == storage
+
+    items = list(
+        child_storage.find(base_dir="any/path", file_pattern=re.compile(r".*\.txt$"))
+    )
+    assert items == [("tests/fixtures/text/input/dulce.txt", {})]