fix tests

deepset-ai · Feb 29, 2024 · 43c9e82 · 43c9e82
1 parent d1d46a5
commit 43c9e82
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 36 deletions.
diff --git a/integrations/chroma/pyproject.toml b/integrations/chroma/pyproject.toml
@@ -189,6 +189,7 @@ module = [
   "chromadb.*",
   "haystack.*",
   "haystack_integrations.*",
-  "pytest.*"
+  "pytest.*",
+  "numpy.*"
 ]
 ignore_missing_imports = true
diff --git a/integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py b/integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py
@@ -6,8 +6,8 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import chromadb
-from chromadb.api.types import GetResult, QueryResult, validate_where, validate_where_document
 import numpy as np
+from chromadb.api.types import GetResult, QueryResult, validate_where, validate_where_document
 from haystack import default_from_dict, default_to_dict
 from haystack.dataclasses import Document
 from haystack.document_stores.types import DuplicatePolicy

diff --git a/integrations/chroma/tests/test_document_store.py b/integrations/chroma/tests/test_document_store.py
@@ -97,12 +97,12 @@ def test_to_json(self, request):
         )
         ds_dict = ds.to_dict()
         assert ds_dict == {
-            'type': 'haystack_integrations.document_stores.chroma.document_store.ChromaDocumentStore',
-            'init_parameters': {
-                'collection_name': 'test_to_json',
-                'embedding_function': 'HuggingFaceEmbeddingFunction',
-                'persist_path': None,
-                'api_key': '1234567890',
+            "type": "haystack_integrations.document_stores.chroma.document_store.ChromaDocumentStore",
+            "init_parameters": {
+                "collection_name": "test_to_json",
+                "embedding_function": "HuggingFaceEmbeddingFunction",
+                "persist_path": None,
+                "api_key": "1234567890",
             },
         }
 
@@ -111,12 +111,12 @@ def test_from_json(self):
         collection_name = "test_collection"
         function_name = "HuggingFaceEmbeddingFunction"
         ds_dict = {
-            'type': 'haystack_integrations.document_stores.chroma.document_store.ChromaDocumentStore',
-            'init_parameters': {
-                'collection_name': 'test_collection',
-                'embedding_function': 'HuggingFaceEmbeddingFunction',
-                'persist_path': None,
-                'api_key': '1234567890',
+            "type": "haystack_integrations.document_stores.chroma.document_store.ChromaDocumentStore",
+            "init_parameters": {
+                "collection_name": "test_collection",
+                "embedding_function": "HuggingFaceEmbeddingFunction",
+                "persist_path": None,
+                "api_key": "1234567890",
             },
         }
 

diff --git a/integrations/chroma/tests/test_retriever.py b/integrations/chroma/tests/test_retriever.py
@@ -13,17 +13,17 @@ def test_retriever_to_json(request):
     )
     retriever = ChromaQueryTextRetriever(ds, filters={"foo": "bar"}, top_k=99)
     assert retriever.to_dict() == {
-        'type': 'haystack_integrations.components.retrievers.chroma.retriever.ChromaQueryTextRetriever',
-        'init_parameters': {
-            'filters': {'foo': 'bar'},
-            'top_k': 99,
-            'document_store': {
-                'type': 'haystack_integrations.document_stores.chroma.document_store.ChromaDocumentStore',
-                'init_parameters': {
-                    'collection_name': 'test_retriever_to_json',
-                    'embedding_function': 'HuggingFaceEmbeddingFunction',
-                    'persist_path': None,
-                    'api_key': '1234567890',
+        "type": "haystack_integrations.components.retrievers.chroma.retriever.ChromaQueryTextRetriever",
+        "init_parameters": {
+            "filters": {"foo": "bar"},
+            "top_k": 99,
+            "document_store": {
+                "type": "haystack_integrations.document_stores.chroma.document_store.ChromaDocumentStore",
+                "init_parameters": {
+                    "collection_name": "test_retriever_to_json",
+                    "embedding_function": "HuggingFaceEmbeddingFunction",
+                    "persist_path": None,
+                    "api_key": "1234567890",
                 },
             },
         },
@@ -33,17 +33,17 @@ def test_retriever_to_json(request):
 @pytest.mark.integration
 def test_retriever_from_json(request):
     data = {
-        'type': 'haystack_integrations.components.retrievers.chroma.retriever.ChromaQueryTextRetriever',
-        'init_parameters': {
-            'filters': {'bar': 'baz'},
-            'top_k': 42,
-            'document_store': {
-                'type': 'haystack_integrations.document_stores.chroma.document_store.ChromaDocumentStore',
-                'init_parameters': {
-                    'collection_name': 'test_retriever_from_json',
-                    'embedding_function': 'HuggingFaceEmbeddingFunction',
-                    'persist_path': '.',
-                    'api_key': '1234567890',
+        "type": "haystack_integrations.components.retrievers.chroma.retriever.ChromaQueryTextRetriever",
+        "init_parameters": {
+            "filters": {"bar": "baz"},
+            "top_k": 42,
+            "document_store": {
+                "type": "haystack_integrations.document_stores.chroma.document_store.ChromaDocumentStore",
+                "init_parameters": {
+                    "collection_name": "test_retriever_from_json",
+                    "embedding_function": "HuggingFaceEmbeddingFunction",
+                    "persist_path": ".",
+                    "api_key": "1234567890",
                 },
             },
         },