diff --git a/integrations/gradient/tests/test_gradient_document_embedder.py b/integrations/gradient/tests/test_gradient_document_embedder.py index f78e19f1f..6e75360fe 100644 --- a/integrations/gradient/tests/test_gradient_document_embedder.py +++ b/integrations/gradient/tests/test_gradient_document_embedder.py @@ -13,7 +13,6 @@ class TestGradientDocumentEmbedder: - @pytest.mark.unit def test_init_from_env(self, monkeypatch): monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", access_token) monkeypatch.setenv("GRADIENT_WORKSPACE_ID", workspace_id) @@ -23,28 +22,24 @@ def test_init_from_env(self, monkeypatch): assert embedder._gradient.workspace_id == workspace_id assert embedder._gradient._api_client.configuration.access_token == access_token - @pytest.mark.unit def test_init_without_access_token(self, monkeypatch): monkeypatch.delenv("GRADIENT_ACCESS_TOKEN", raising=False) with pytest.raises(ValueError): GradientDocumentEmbedder(workspace_id=workspace_id) - @pytest.mark.unit def test_init_without_workspace(self, monkeypatch): monkeypatch.delenv("GRADIENT_WORKSPACE_ID", raising=False) with pytest.raises(ValueError): GradientDocumentEmbedder(access_token=access_token) - @pytest.mark.unit def test_init_from_params(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) assert embedder is not None assert embedder._gradient.workspace_id == workspace_id assert embedder._gradient._api_client.configuration.access_token == access_token - @pytest.mark.unit def test_init_from_params_precedence(self, monkeypatch): monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", "env_access_token") monkeypatch.setenv("GRADIENT_WORKSPACE_ID", "env_workspace_id") @@ -54,7 +49,6 @@ def test_init_from_params_precedence(self, monkeypatch): assert embedder._gradient.workspace_id == workspace_id assert embedder._gradient._api_client.configuration.access_token == access_token - @pytest.mark.unit def test_to_dict(self): component = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) data = component.to_dict() @@ -63,14 +57,12 @@ def test_to_dict(self): "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, } - @pytest.mark.unit def test_warmup(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._gradient.get_embeddings_model = MagicMock() embedder.warm_up() embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") - @pytest.mark.unit def test_warmup_doesnt_reload(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._gradient.get_embeddings_model = MagicMock(default_return_value="fake model") @@ -78,14 +70,12 @@ def test_warmup_doesnt_reload(self): embedder.warm_up() embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") - @pytest.mark.unit def test_run_fail_if_not_warmed_up(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) with pytest.raises(RuntimeError, match="warm_up()"): embedder.run(documents=[Document(content=f"document number {i}") for i in range(5)]) - @pytest.mark.unit def test_run(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() @@ -105,7 +95,6 @@ def test_run(self): assert isinstance(doc.embedding, list) assert isinstance(doc.embedding[0], float) - @pytest.mark.unit def test_run_batch(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() @@ -126,7 +115,6 @@ def test_run_batch(self): assert isinstance(doc.embedding, list) assert isinstance(doc.embedding[0], float) - @pytest.mark.unit def test_run_custom_batch(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id, batch_size=20) embedder._embedding_model = NonCallableMagicMock() @@ -148,7 +136,6 @@ def test_run_custom_batch(self): assert isinstance(doc.embedding, list) assert isinstance(doc.embedding[0], float) - @pytest.mark.unit def test_run_empty(self): embedder = GradientDocumentEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() diff --git a/integrations/gradient/tests/test_gradient_rag_pipelines.py b/integrations/gradient/tests/test_gradient_rag_pipelines.py index c5ec4affb..9d3b11486 100644 --- a/integrations/gradient/tests/test_gradient_rag_pipelines.py +++ b/integrations/gradient/tests/test_gradient_rag_pipelines.py @@ -5,9 +5,9 @@ from haystack import Document, Pipeline from haystack.components.builders.answer_builder import AnswerBuilder from haystack.components.builders.prompt_builder import PromptBuilder -from haystack.components.retrievers import InMemoryEmbeddingRetriever +from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever from haystack.components.writers import DocumentWriter -from haystack.document_stores import InMemoryDocumentStore +from haystack.document_stores.in_memory import InMemoryDocumentStore from gradient_haystack.embedders.gradient_document_embedder import GradientDocumentEmbedder from gradient_haystack.embedders.gradient_text_embedder import GradientTextEmbedder diff --git a/integrations/gradient/tests/test_gradient_text_embedder.py b/integrations/gradient/tests/test_gradient_text_embedder.py index 98cb1d0fb..7ae846e93 100644 --- a/integrations/gradient/tests/test_gradient_text_embedder.py +++ b/integrations/gradient/tests/test_gradient_text_embedder.py @@ -12,7 +12,6 @@ class TestGradientTextEmbedder: - @pytest.mark.unit def test_init_from_env(self, monkeypatch): monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", access_token) monkeypatch.setenv("GRADIENT_WORKSPACE_ID", workspace_id) @@ -22,28 +21,24 @@ def test_init_from_env(self, monkeypatch): assert embedder._gradient.workspace_id == workspace_id assert embedder._gradient._api_client.configuration.access_token == access_token - @pytest.mark.unit def test_init_without_access_token(self, monkeypatch): monkeypatch.delenv("GRADIENT_ACCESS_TOKEN", raising=False) with pytest.raises(ValueError): GradientTextEmbedder(workspace_id=workspace_id) - @pytest.mark.unit def test_init_without_workspace(self, monkeypatch): monkeypatch.delenv("GRADIENT_WORKSPACE_ID", raising=False) with pytest.raises(ValueError): GradientTextEmbedder(access_token=access_token) - @pytest.mark.unit def test_init_from_params(self): embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) assert embedder is not None assert embedder._gradient.workspace_id == workspace_id assert embedder._gradient._api_client.configuration.access_token == access_token - @pytest.mark.unit def test_init_from_params_precedence(self, monkeypatch): monkeypatch.setenv("GRADIENT_ACCESS_TOKEN", "env_access_token") monkeypatch.setenv("GRADIENT_WORKSPACE_ID", "env_workspace_id") @@ -53,7 +48,6 @@ def test_init_from_params_precedence(self, monkeypatch): assert embedder._gradient.workspace_id == workspace_id assert embedder._gradient._api_client.configuration.access_token == access_token - @pytest.mark.unit def test_to_dict(self): component = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) data = component.to_dict() @@ -62,14 +56,12 @@ def test_to_dict(self): "init_parameters": {"workspace_id": workspace_id, "model_name": "bge-large"}, } - @pytest.mark.unit def test_warmup(self): embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._gradient.get_embeddings_model = MagicMock() embedder.warm_up() embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") - @pytest.mark.unit def test_warmup_doesnt_reload(self): embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._gradient.get_embeddings_model = MagicMock(default_return_value="fake model") @@ -77,14 +69,12 @@ def test_warmup_doesnt_reload(self): embedder.warm_up() embedder._gradient.get_embeddings_model.assert_called_once_with(slug="bge-large") - @pytest.mark.unit def test_run_fail_if_not_warmed_up(self): embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) with pytest.raises(RuntimeError, match="warm_up()"): embedder.run(text="The food was delicious") - @pytest.mark.unit def test_run_fail_when_no_embeddings_returned(self): embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() @@ -94,7 +84,6 @@ def test_run_fail_when_no_embeddings_returned(self): _result = embedder.run(text="The food was delicious") embedder._embedding_model.embed.assert_called_once_with(inputs=[{"input": "The food was delicious"}]) - @pytest.mark.unit def test_run_empty_string(self): embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock() @@ -108,7 +97,6 @@ def test_run_empty_string(self): assert len(result["embedding"]) == 1024 # 1024 is the bge-large embedding size assert all(isinstance(x, float) for x in result["embedding"]) - @pytest.mark.unit def test_run(self): embedder = GradientTextEmbedder(access_token=access_token, workspace_id=workspace_id) embedder._embedding_model = NonCallableMagicMock()