Add Store class factory (#5530)

* Add Store class factory

* Add release notes
This commit is contained in:
Silvano Cerza 2023-08-09 13:09:36 +02:00 committed by GitHub
parent ff86af576a
commit 83fce1bd72
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 179 additions and 0 deletions

View File

View File

@ -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)

View File

@ -0,0 +1,4 @@
---
features:
- |
Add utility function `store_class` factory to create `Store`s for testing purposes.

View File

@ -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