From 53b77dda6c66b0bf21de4ea4983872667c9667e4 Mon Sep 17 00:00:00 2001
From: Silvano Cerza <3314350+silvanocerza@users.noreply.github.com>
Date: Fri, 17 Nov 2023 19:25:16 +0100
Subject: [PATCH] Move tests for write_documents from DocumentStoreBaseTests to
 separate class (#6334)

---
 haystack/preview/testing/document_store.py | 102 ++++++++++++---------
 1 file changed, 59 insertions(+), 43 deletions(-)

diff --git a/haystack/preview/testing/document_store.py b/haystack/preview/testing/document_store.py
index 08ab974ea..d353611ca 100644
--- a/haystack/preview/testing/document_store.py
+++ b/haystack/preview/testing/document_store.py
@@ -43,7 +43,65 @@ class CountDocumentsTest:
         assert docstore.count_documents() == 3
 
 
-class DocumentStoreBaseTests(CountDocumentsTest):
+class WriteDocumentsTest:
+    """
+    Utility class to test a Document Store `write_documents` method.
+
+    To use it create a custom test class and override the `docstore` fixture to return your Document Store.
+    Example usage:
+
+    ```python
+    class MyDocumentStoreTest(WriteDocumentsTest):
+        @pytest.fixture
+        def docstore(self):
+            return MyDocumentStore()
+    ```
+    """
+
+    @pytest.mark.unit
+    def test_write(self, docstore: DocumentStore):
+        doc = Document(content="test doc")
+        docstore.write_documents([doc])
+        assert docstore.filter_documents(filters={"id": doc.id}) == [doc]
+
+    @pytest.mark.unit
+    def test_write_duplicate_fail(self, docstore: DocumentStore):
+        doc = Document(content="test doc")
+        docstore.write_documents([doc])
+        with pytest.raises(DuplicateDocumentError, match=f"ID '{doc.id}' already exists."):
+            docstore.write_documents(documents=[doc], policy=DuplicatePolicy.FAIL)
+        assert docstore.filter_documents(filters={"id": doc.id}) == [doc]
+
+    @pytest.mark.unit
+    def test_write_duplicate_skip(self, docstore: DocumentStore):
+        doc = Document(content="test doc")
+        docstore.write_documents([doc])
+        docstore.write_documents(documents=[doc], policy=DuplicatePolicy.SKIP)
+        assert docstore.filter_documents(filters={"id": doc.id}) == [doc]
+
+    @pytest.mark.unit
+    def test_write_duplicate_overwrite(self, docstore: DocumentStore):
+        doc1 = Document(content="test doc 1")
+        doc2 = Document(content="test doc 2")
+        object.__setattr__(doc2, "id", doc1.id)  # Make two docs with different content but same ID
+
+        docstore.write_documents([doc2])
+        assert docstore.filter_documents(filters={"id": doc1.id}) == [doc2]
+        docstore.write_documents(documents=[doc1], policy=DuplicatePolicy.OVERWRITE)
+        assert docstore.filter_documents(filters={"id": doc1.id}) == [doc1]
+
+    @pytest.mark.unit
+    def test_write_not_docs(self, docstore: DocumentStore):
+        with pytest.raises(ValueError):
+            docstore.write_documents(["not a document for sure"])  # type: ignore
+
+    @pytest.mark.unit
+    def test_write_not_list(self, docstore: DocumentStore):
+        with pytest.raises(ValueError):
+            docstore.write_documents("not a list actually")  # type: ignore
+
+
+class DocumentStoreBaseTests(CountDocumentsTest, WriteDocumentsTest):
     @pytest.fixture
     def docstore(self) -> DocumentStore:
         raise NotImplementedError()
@@ -559,48 +617,6 @@ class DocumentStoreBaseTests(CountDocumentsTest):
             )
         ]
 
-    @pytest.mark.unit
-    def test_write(self, docstore: DocumentStore):
-        doc = Document(content="test doc")
-        docstore.write_documents([doc])
-        assert docstore.filter_documents(filters={"id": doc.id}) == [doc]
-
-    @pytest.mark.unit
-    def test_write_duplicate_fail(self, docstore: DocumentStore):
-        doc = Document(content="test doc")
-        docstore.write_documents([doc])
-        with pytest.raises(DuplicateDocumentError, match=f"ID '{doc.id}' already exists."):
-            docstore.write_documents(documents=[doc], policy=DuplicatePolicy.FAIL)
-        assert docstore.filter_documents(filters={"id": doc.id}) == [doc]
-
-    @pytest.mark.unit
-    def test_write_duplicate_skip(self, docstore: DocumentStore):
-        doc = Document(content="test doc")
-        docstore.write_documents([doc])
-        docstore.write_documents(documents=[doc], policy=DuplicatePolicy.SKIP)
-        assert docstore.filter_documents(filters={"id": doc.id}) == [doc]
-
-    @pytest.mark.unit
-    def test_write_duplicate_overwrite(self, docstore: DocumentStore):
-        doc1 = Document(content="test doc 1")
-        doc2 = Document(content="test doc 2")
-        object.__setattr__(doc2, "id", doc1.id)  # Make two docs with different content but same ID
-
-        docstore.write_documents([doc2])
-        assert docstore.filter_documents(filters={"id": doc1.id}) == [doc2]
-        docstore.write_documents(documents=[doc1], policy=DuplicatePolicy.OVERWRITE)
-        assert docstore.filter_documents(filters={"id": doc1.id}) == [doc1]
-
-    @pytest.mark.unit
-    def test_write_not_docs(self, docstore: DocumentStore):
-        with pytest.raises(ValueError):
-            docstore.write_documents(["not a document for sure"])  # type: ignore
-
-    @pytest.mark.unit
-    def test_write_not_list(self, docstore: DocumentStore):
-        with pytest.raises(ValueError):
-            docstore.write_documents("not a list actually")  # type: ignore
-
     @pytest.mark.unit
     def test_delete_empty(self, docstore: DocumentStore):
         with pytest.raises(MissingDocumentError):