fix test

Author: dayesouza

graphrag/index/input/factory.py

@@ -81,8 +81,6 @@ async def create_input(
             # Collapse the metadata columns into a single JSON object column
             result["metadata"] = result[config.metadata].to_dict(orient="records")
 
-            # Drop the original metadata columns after collapsing them
-            result.drop(columns=config.metadata, inplace=True)
         return cast("pd.DataFrame", result)
 
     msg = f"Unknown input type {config.file_type}"

graphrag/storage/file_pipeline_storage.py

@@ -178,4 +178,4 @@ def _create_progress_status(
         total_items=num_total,
         completed_items=num_loaded + num_filtered,
         description=f"{num_loaded} files loaded ({num_filtered} filtered)",
-    )
+    )

graphrag/storage/pipeline_storage.py

@@ -6,6 +6,7 @@
 import re
 from abc import ABCMeta, abstractmethod
 from collections.abc import Iterator
+from datetime import datetime
 from typing import Any
 
 from graphrag.logger.base import ProgressLogger
@@ -92,3 +93,10 @@ def get_creation_date(self, key: str) -> str:
         -------
             - output - The creation date for the given key.
         """
+
+
+def get_timestamp_formatted_with_local_tz(timestamp: datetime) -> str:
+    """Get the formatted timestamp with the local time zone."""
+    creation_time_local = timestamp.astimezone()
+
+    return creation_time_local.strftime("%Y-%m-%d %H:%M:%S %z")

tests/integration/storage/test_cosmosdb_storage.py

@@ -109,6 +109,7 @@ async def test_clear():
     finally:
         await storage.clear()
 
+
 async def test_get_creation_date():
     storage = CosmosDBPipelineStorage(
         connection_string=WELL_KNOWN_COSMOS_CONNECTION_STRING,

tests/verbs/data/source_documents.parquet

@@ -37,7 +37,7 @@ async def test_create_final_documents():
     compare_outputs(actual, expected)
 
 
-async def test_create_final_documents_with_attribute_columns():
+async def test_create_final_documents_with_metadata_column():
     expected = load_test_table(workflow_name)
 
     context = await create_test_context(
@@ -61,6 +61,6 @@ async def test_create_final_documents_with_attribute_columns():
     compare_outputs(
         actual, expected, columns=["id", "human_readable_id", "text", "text_unit_ids"]
     )
-    assert len(actual.columns) == 5
-    assert "title" not in actual.columns
+    assert len(actual.columns) == 6
+    assert "title" in actual.columns
     assert "metadata" in actual.columns