fix integration tests (#450)

deepset-ai · Feb 20, 2024 · 3a6cf81 · 3a6cf81
1 parent f3ea6be
commit 3a6cf81
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/integrations/astra/src/haystack_integrations/document_stores/astra/astra_client.py b/integrations/astra/src/haystack_integrations/document_stores/astra/astra_client.py
@@ -184,9 +184,9 @@ def _query(self, vector, top_k, filters=None):
 
     def find_documents(self, find_query):
         response_dict = self._astra_db_collection.find(
-            filter=find_query["filter"],
-            projection=find_query["sort"],
-            options=find_query["options"],
+            filter=find_query.get("filter"),
+            projection=find_query.get("sort"),
+            options=find_query.get("options"),
         )
 
         if "data" in response_dict and "documents" in response_dict["data"]:

diff --git a/integrations/astra/src/haystack_integrations/document_stores/astra/document_store.py b/integrations/astra/src/haystack_integrations/document_stores/astra/document_store.py
@@ -161,7 +161,8 @@ def _convert_input_document(document: Union[dict, Document]):
 
             if "dataframe" in document_dict and document_dict["dataframe"] is not None:
                 document_dict["dataframe"] = document_dict.pop("dataframe").to_json()
-            document_dict["$vector"] = document_dict.pop("embedding", None)
+            if embedding := document_dict.pop("embedding", []):
+                document_dict["$vector"] = embedding
 
             return document_dict
 

diff --git a/integrations/astra/tests/test_document_store.py b/integrations/astra/tests/test_document_store.py
@@ -29,7 +29,7 @@ def document_store(self) -> AstraDocumentStore:
         return AstraDocumentStore(
             collection_name="haystack_integration",
             duplicates_policy=DuplicatePolicy.OVERWRITE,
-            embedding_dim=768,
+            embedding_dimension=768,
         )
 
     @pytest.fixture(autouse=True)