From 43c9e82dea6a1d33d2ae6757fcc8d0f0b8802670 Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Thu, 29 Feb 2024 17:41:46 +0100 Subject: [PATCH] fix tests --- integrations/chroma/pyproject.toml | 3 +- .../document_stores/chroma/document_store.py | 2 +- .../chroma/tests/test_document_store.py | 24 +++++----- integrations/chroma/tests/test_retriever.py | 44 +++++++++---------- 4 files changed, 37 insertions(+), 36 deletions(-) diff --git a/integrations/chroma/pyproject.toml b/integrations/chroma/pyproject.toml index 2653c491f..f8265a1e7 100644 --- 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 index e05e9b636..e81a94ef0 100644 --- 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 index d4a078400..8d61e63ed 100644 --- 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 index 75a30f242..b430e5fda 100644 --- 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", }, }, },