test: Opensearch - add unit tests (#3182)

Author: anakin87

integrations/opensearch/src/haystack_integrations/components/retrievers/opensearch/bm25_retriever.py

@@ -13,6 +13,8 @@
 
 from haystack_integrations.document_stores.opensearch import OpenSearchDocumentStore
 
+from .utils import _resolve_document_store
+
 logger = logging.getLogger(__name__)
 
 
@@ -268,13 +270,7 @@ def run(
             custom_query=custom_query,
         )
 
-        if document_store is not None:
-            if not isinstance(document_store, OpenSearchDocumentStore):
-                msg = "document_store must be an instance of OpenSearchDocumentStore"
-                raise ValueError(msg)
-            doc_store = document_store
-        else:
-            doc_store = self._document_store
+        doc_store = _resolve_document_store(document_store, self._document_store)
 
         try:
             docs = doc_store._bm25_retrieval(**bm25_args)  # example for BM25Retriever
@@ -335,13 +331,7 @@ async def run_async(  # pylint: disable=too-many-positional-arguments
             custom_query=custom_query,
         )
 
-        if document_store is not None:
-            if not isinstance(document_store, OpenSearchDocumentStore):
-                msg = "document_store must be an instance of OpenSearchDocumentStore"
-                raise ValueError(msg)
-            doc_store = document_store
-        else:
-            doc_store = self._document_store
+        doc_store = _resolve_document_store(document_store, self._document_store)
 
         try:
             docs = await doc_store._bm25_retrieval_async(**bm25_args)

integrations/opensearch/src/haystack_integrations/components/retrievers/opensearch/embedding_retriever.py

@@ -13,6 +13,8 @@
 
 from haystack_integrations.document_stores.opensearch import OpenSearchDocumentStore
 
+from .utils import _resolve_document_store
+
 logger = logging.getLogger(__name__)
 
 
@@ -257,13 +259,7 @@ def run(
 
         docs: list[Document] = []
 
-        if document_store is not None:
-            if not isinstance(document_store, OpenSearchDocumentStore):
-                msg = "document_store must be an instance of OpenSearchDocumentStore"
-                raise ValueError(msg)
-            doc_store = document_store
-        else:
-            doc_store = self._document_store
+        doc_store = _resolve_document_store(document_store, self._document_store)
 
         try:
             docs = doc_store._embedding_retrieval(
@@ -383,13 +379,7 @@ async def run_async(
 
         docs: list[Document] = []
 
-        if document_store is not None:
-            if not isinstance(document_store, OpenSearchDocumentStore):
-                msg = "document_store must be an instance of OpenSearchDocumentStore"
-                raise ValueError(msg)
-            doc_store = document_store
-        else:
-            doc_store = self._document_store
+        doc_store = _resolve_document_store(document_store, self._document_store)
 
         try:
             docs = await doc_store._embedding_retrieval_async(

integrations/opensearch/src/haystack_integrations/components/retrievers/opensearch/metadata_retriever.py

@@ -9,6 +9,8 @@
 
 from haystack_integrations.document_stores.opensearch import OpenSearchDocumentStore
 
+from .utils import _resolve_document_store
+
 logger = logging.getLogger(__name__)
 
 
@@ -252,10 +254,7 @@ def run(
             # Returns: {"metadata": [{"category": "Python", "status": "active", "priority": 1}]}
             ```
         """
-        doc_store = document_store or self._document_store
-        if not isinstance(doc_store, OpenSearchDocumentStore):
-            msg = "document_store must be an instance of OpenSearchDocumentStore"
-            raise ValueError(msg)
+        doc_store = _resolve_document_store(document_store, self._document_store)
 
         fields_to_use = metadata_fields if metadata_fields is not None else self._metadata_fields
         top_k_to_use = top_k if top_k is not None else self._top_k
@@ -366,10 +365,7 @@ async def run_async(
             # Returns: {"metadata": [{"category": "Python", "status": "active", "priority": 1}]}
             ```
         """
-        doc_store = document_store or self._document_store
-        if not isinstance(doc_store, OpenSearchDocumentStore):
-            msg = "document_store must be an instance of OpenSearchDocumentStore"
-            raise ValueError(msg)
+        doc_store = _resolve_document_store(document_store, self._document_store)
 
         fields_to_use = metadata_fields if metadata_fields is not None else self._metadata_fields
         top_k_to_use = top_k if top_k is not None else self._top_k

integrations/opensearch/src/haystack_integrations/components/retrievers/opensearch/sql_retriever.py

@@ -8,6 +8,8 @@
 
 from haystack_integrations.document_stores.opensearch import OpenSearchDocumentStore
 
+from .utils import _resolve_document_store
+
 logger = logging.getLogger(__name__)
 
 
@@ -108,13 +110,7 @@ def run(
             # For aggregate queries: result["result"]["aggregations"] contains aggregations
             ```
         """
-        if document_store is not None:
-            if not isinstance(document_store, OpenSearchDocumentStore):
-                msg = "document_store must be an instance of OpenSearchDocumentStore"
-                raise ValueError(msg)
-            doc_store = document_store
-        else:
-            doc_store = self._document_store
+        doc_store = _resolve_document_store(document_store, self._document_store)
 
         fetch_size = fetch_size if fetch_size is not None else self._fetch_size
 
@@ -163,13 +159,7 @@ async def run_async(
             # For aggregate queries: result["result"]["aggregations"] contains aggregations
             ```
         """
-        if document_store is not None:
-            if not isinstance(document_store, OpenSearchDocumentStore):
-                msg = "document_store must be an instance of OpenSearchDocumentStore"
-                raise ValueError(msg)
-            doc_store = document_store
-        else:
-            doc_store = self._document_store
+        doc_store = _resolve_document_store(document_store, self._document_store)
 
         fetch_size = fetch_size if fetch_size is not None else self._fetch_size
 

integrations/opensearch/src/haystack_integrations/components/retrievers/opensearch/utils.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+from haystack_integrations.document_stores.opensearch import OpenSearchDocumentStore
+
+
+def _resolve_document_store(
+    runtime_document_store: OpenSearchDocumentStore | None,
+    default_document_store: OpenSearchDocumentStore,
+) -> OpenSearchDocumentStore:
+    """
+    Return the runtime document store if provided and valid, otherwise the default one.
+
+    :raises ValueError: If `runtime_document_store` is not None and not an OpenSearchDocumentStore.
+    """
+    if runtime_document_store is None:
+        return default_document_store
+    if not isinstance(runtime_document_store, OpenSearchDocumentStore):
+        msg = "document_store must be an instance of OpenSearchDocumentStore"
+        raise ValueError(msg)
+    return runtime_document_store