From 83fce1bd72d59b9b87165e1d79fdf02d0184f645 Mon Sep 17 00:00:00 2001
From: Silvano Cerza <3314350+silvanocerza@users.noreply.github.com>
Date: Wed, 9 Aug 2023 13:09:36 +0200
Subject: [PATCH] Add Store class factory (#5530)

* Add Store class factory

* Add release notes
---
 haystack/preview/testing/__init__.py          |   0
 haystack/preview/testing/factory.py           | 114 ++++++++++++++++++
 .../notes/store-factory-91e7da46aeb7ff21.yaml |   4 +
 test/preview/testing/test_factory.py          |  61 ++++++++++
 4 files changed, 179 insertions(+)
 create mode 100644 haystack/preview/testing/__init__.py
 create mode 100644 haystack/preview/testing/factory.py
 create mode 100644 releasenotes/notes/store-factory-91e7da46aeb7ff21.yaml
 create mode 100644 test/preview/testing/test_factory.py

diff --git a/haystack/preview/testing/__init__.py b/haystack/preview/testing/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/haystack/preview/testing/factory.py b/haystack/preview/testing/factory.py
new file mode 100644
index 000000000..baeb02844
--- /dev/null
+++ b/haystack/preview/testing/factory.py
@@ -0,0 +1,114 @@
+from typing import Any, Dict, Optional, Tuple, Type, List
+
+from haystack.preview.dataclasses import Document
+from haystack.preview.document_stores import store, Store, DuplicatePolicy
+
+
+def store_class(
+    name: str,
+    documents: Optional[List[Document]] = None,
+    documents_count: Optional[int] = None,
+    bases: Optional[Tuple[type, ...]] = None,
+    extra_fields: Optional[Dict[str, Any]] = None,
+) -> Type[Store]:
+    """
+    Utility function to create a Store class with the given name and list of documents.
+
+    If `documents` is set but `documents_count` is not, `documents_count` will be the length
+    of `documents`.
+    If both are set explicitly they don't influence each other.
+
+    `write_documents()` and `delete_documents()` are no-op.
+    You can override them using `extra_fields`.
+
+    ### Usage
+
+    Create a store class that returns no documents:
+    ```python
+    MyFakeStore = store_class("MyFakeComponent")
+    store = MyFakeStore()
+    assert store.documents_count() == 0
+    assert store.filter_documents() == []
+    ```
+
+    Create a store class that returns a single document:
+    ```python
+    doc = Document(id="fake_id", content="Fake content")
+    MyFakeStore = store_class("MyFakeComponent", documents=[doc])
+    store = MyFakeStore()
+    assert store.documents_count() == 1
+    assert store.filter_documents() == [doc]
+    ```
+
+    Create a store class that returns no document but returns a custom count:
+    ```python
+    MyFakeStore = store_class("MyFakeComponent", documents_count=100)
+    store = MyFakeStore()
+    assert store.documents_count() == 100
+    assert store.filter_documents() == []
+    ```
+
+    Create a store class that returns a document and a custom count:
+    ```python
+    doc = Document(id="fake_id", content="Fake content")
+    MyFakeStore = store_class("MyFakeComponent", documents=[doc], documents_count=100)
+    store = MyFakeStore()
+    assert store.documents_count() == 100
+    assert store.filter_documents() == [doc]
+    ```
+
+    Create a store class with a custom base class:
+    ```python
+    MyFakeStore = store_class(
+        "MyFakeStore",
+        bases=(MyBaseClass,)
+    )
+    store = MyFakeStore()
+    assert isinstance(store, MyBaseClass)
+    ```
+
+    Create a store class with an extra field `my_field`:
+    ```python
+    MyFakeStore = store_class(
+        "MyFakeStore",
+        extra_fields={"my_field": 10}
+    )
+    store = MyFakeStore()
+    assert store.my_field == 10
+    ```
+    """
+
+    if documents is not None and documents_count is None:
+        documents_count = len(documents)
+    elif documents_count is None:
+        documents_count = 0
+
+    def count_documents(self) -> int:
+        return documents_count
+
+    def filter_documents(self, filters: Optional[Dict[str, Any]] = None) -> List[Document]:
+        if documents is not None:
+            return documents
+        return []
+
+    def write_documents(self, documents: List[Document], policy: DuplicatePolicy = DuplicatePolicy.FAIL) -> None:
+        return
+
+    def delete_documents(self, document_ids: List[str]) -> None:
+        return
+
+    fields = {
+        "count_documents": count_documents,
+        "filter_documents": filter_documents,
+        "write_documents": write_documents,
+        "delete_documents": delete_documents,
+    }
+
+    if extra_fields is not None:
+        fields = {**fields, **extra_fields}
+
+    if bases is None:
+        bases = (object,)
+
+    cls = type(name, bases, fields)
+    return store(cls)
diff --git a/releasenotes/notes/store-factory-91e7da46aeb7ff21.yaml b/releasenotes/notes/store-factory-91e7da46aeb7ff21.yaml
new file mode 100644
index 000000000..2330a93aa
--- /dev/null
+++ b/releasenotes/notes/store-factory-91e7da46aeb7ff21.yaml
@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Add utility function `store_class` factory to create `Store`s for testing purposes.
diff --git a/test/preview/testing/test_factory.py b/test/preview/testing/test_factory.py
new file mode 100644
index 000000000..2e6626c53
--- /dev/null
+++ b/test/preview/testing/test_factory.py
@@ -0,0 +1,61 @@
+import pytest
+
+from haystack.preview.dataclasses import Document
+from haystack.preview.testing.factory import store_class
+from haystack.preview.document_stores.decorator import store
+
+
+@pytest.mark.unit
+def test_store_class_default():
+    MyStore = store_class("MyStore")
+    store = MyStore()
+    assert store.count_documents() == 0
+    assert store.filter_documents() == []
+    assert store.write_documents([]) is None
+    assert store.delete_documents([]) is None
+
+
+@pytest.mark.unit
+def test_store_class_is_registered():
+    MyStore = store_class("MyStore")
+    assert store.registry["MyStore"] == MyStore
+
+
+@pytest.mark.unit
+def test_store_class_with_documents():
+    doc = Document(id="fake_id", content="This is a document")
+    MyStore = store_class("MyStore", documents=[doc])
+    store = MyStore()
+    assert store.count_documents() == 1
+    assert store.filter_documents() == [doc]
+
+
+@pytest.mark.unit
+def test_store_class_with_documents_count():
+    MyStore = store_class("MyStore", documents_count=100)
+    store = MyStore()
+    assert store.count_documents() == 100
+    assert store.filter_documents() == []
+
+
+@pytest.mark.unit
+def test_store_class_with_documents_and_documents_count():
+    doc = Document(id="fake_id", content="This is a document")
+    MyStore = store_class("MyStore", documents=[doc], documents_count=100)
+    store = MyStore()
+    assert store.count_documents() == 100
+    assert store.filter_documents() == [doc]
+
+
+@pytest.mark.unit
+def test_store_class_with_bases():
+    MyStore = store_class("MyStore", bases=(Exception,))
+    store = MyStore()
+    assert isinstance(store, Exception)
+
+
+@pytest.mark.unit
+def test_store_class_with_extra_fields():
+    MyStore = store_class("MyStore", extra_fields={"my_field": 10})
+    store = MyStore()
+    assert store.my_field == 10