use factory

deepset-ai · Aug 29, 2023 · f2013b5 · f2013b5
1 parent 9062c07
commit f2013b5
Showing 1 changed file with 10 additions and 30 deletions.
diff --git a/test/preview/components/writers/test_document_writer.py b/test/preview/components/writers/test_document_writer.py
@@ -3,70 +3,50 @@
 import pytest
 
 from haystack.preview import Document, DeserializationError
-from haystack.preview.document_stores import document_store
+from haystack.preview.testing.factory import document_store_class
 from haystack.preview.components.writers.document_writer import DocumentWriter
 from haystack.preview.document_stores import DuplicatePolicy
 
 
 class TestDocumentWriter:
     @pytest.mark.unit
     def test_to_dict(self):
-        mocked_document_store = MagicMock()
-        mocked_document_store.to_dict.return_value = {
-            "type": "MockedDocumentStore",
-            "init_parameters": {"parameter": 100},
-        }
-        component = DocumentWriter(document_store=mocked_document_store)
+        mocked_docstore_class = document_store_class("MockedDocumentStore")
+        component = DocumentWriter(document_store=mocked_docstore_class())
         data = component.to_dict()
         assert data == {
             "type": "DocumentWriter",
             "init_parameters": {
-                "document_store": {"type": "MockedDocumentStore", "init_parameters": {"parameter": 100}},
+                "document_store": {"type": "MockedDocumentStore", "init_parameters": {}},
                 "policy": "FAIL",
             },
         }
 
     @pytest.mark.unit
     def test_to_dict_with_custom_init_parameters(self):
-        mocked_document_store = MagicMock()
-        mocked_document_store.to_dict.return_value = {
-            "type": "MockedDocumentStore",
-            "init_parameters": {"parameter": 100},
-        }
-        component = DocumentWriter(document_store=mocked_document_store, policy=DuplicatePolicy.SKIP)
+        mocked_docstore_class = document_store_class("MockedDocumentStore")
+        component = DocumentWriter(document_store=mocked_docstore_class(), policy=DuplicatePolicy.SKIP)
         data = component.to_dict()
         assert data == {
             "type": "DocumentWriter",
             "init_parameters": {
-                "document_store": {"type": "MockedDocumentStore", "init_parameters": {"parameter": 100}},
+                "document_store": {"type": "MockedDocumentStore", "init_parameters": {}},
                 "policy": "SKIP",
             },
         }
 
     @pytest.mark.unit
     def test_from_dict(self):
-        @document_store
-        class MockedDocumentStore:
-            def __init__(self, parameter):
-                self.parameter = parameter
-
-            def to_dict(self):
-                return {"type": "MockedDocumentStore", "init_parameters": {"parameter": self.parameter}}
-
-            @classmethod
-            def from_dict(cls, data):
-                return cls(parameter=data["init_parameters"]["parameter"])
-
+        mocked_docstore_class = document_store_class("MockedDocumentStore")
         data = {
             "type": "DocumentWriter",
             "init_parameters": {
-                "document_store": {"type": "MockedDocumentStore", "init_parameters": {"parameter": 100}},
+                "document_store": {"type": "MockedDocumentStore", "init_parameters": {}},
                 "policy": "SKIP",
             },
         }
         component = DocumentWriter.from_dict(data)
-        assert isinstance(component.document_store, MockedDocumentStore)
-        assert component.document_store.parameter == 100
+        assert isinstance(component.document_store, mocked_docstore_class)
         assert component.policy == DuplicatePolicy.SKIP
 
     @pytest.mark.unit