Fix: failing test

deepset-ai · Jun 18, 2024 · 31ddcd3 · 31ddcd3
1 parent 24259c4
commit 31ddcd3
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 11 deletions.
diff --git a/haystack/document_stores/in_memory/document_store.py b/haystack/document_stores/in_memory/document_store.py
@@ -364,12 +364,16 @@ def load_from_disk(cls, path: str) -> "InMemoryDocumentStore":
             try:
                 with open(path, "r") as f:
                     data = json.load(f)
-                documents = data.pop("documents")
-                cls_object = default_from_dict(cls, data)
-                cls_object.write_documents(documents=[Document(**doc) for doc in documents])
-                return cls_object
             except Exception as e:
                 raise Exception(f"Error loading InMemoryDocumentStore from disk. error: {e}")
+
+            documents = data.pop("documents")
+            cls_object = default_from_dict(cls, data)
+            cls_object.write_documents(
+                documents=[Document(**doc) for doc in documents], policy=DuplicatePolicy.OVERWRITE
+            )
+            return cls_object
+
         else:
             raise FileNotFoundError(f"File {path} not found.")
 

diff --git a/test/document_stores/test_in_memory.py b/test/document_stores/test_in_memory.py
@@ -19,6 +19,11 @@ class TestMemoryDocumentStore(DocumentStoreBaseTests):  # pylint: disable=R0904
     Test InMemoryDocumentStore's specific features
     """
 
+    @pytest.fixture
+    def tmp_dir(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            yield tmp_dir
+
     @pytest.fixture
     def document_store(self) -> InMemoryDocumentStore:
         return InMemoryDocumentStore(bm25_algorithm="BM25L")
@@ -75,18 +80,16 @@ def test_from_dict(self, mock_regex):
         assert store.bm25_parameters == {"key": "value"}
         assert store.index == "my_cool_index"
 
-    def test_save_to_disk_and_load_from_disk(self, document_store: InMemoryDocumentStore):
+    def test_save_to_disk_and_load_from_disk(self, tmp_dir: str):
         docs = [Document(content="Hello world"), Document(content="Haystack supports multiple languages")]
         document_store = InMemoryDocumentStore()
         document_store.write_documents(docs)
-
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            tmp_dir = tmp_dir + "/tmp.json"
-            document_store.save_to_disk(tmp_dir)
-            document_store_loaded = InMemoryDocumentStore.load_from_disk(tmp_dir)
+        tmp_dir = tmp_dir + "/document_store.json"
+        document_store.save_to_disk(tmp_dir)
+        document_store_loaded = InMemoryDocumentStore.load_from_disk(tmp_dir)
 
         assert document_store_loaded.count_documents() == 2
-        assert document_store_loaded.storage.values() == docs
+        assert list(document_store_loaded.storage.values()) == docs
         assert document_store_loaded.to_dict() == document_store.to_dict()
 
     def test_invalid_bm25_algorithm(self):