From 188708281af3d6b8ccb64bd1e7850c56ecb57721 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 20 Aug 2024 16:49:12 -0300 Subject: [PATCH] test(astra): Refactor AstraDB imports in test_astra_component (#3413) * test(astra): Refactor AstraDB imports in test_astra_component * fix: Add AstraVectorizeComponent to AstraDB tests. --- .../integration/astra/test_astra_component.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/backend/tests/integration/astra/test_astra_component.py b/src/backend/tests/integration/astra/test_astra_component.py index 8dd7c6965..e955539b3 100644 --- a/src/backend/tests/integration/astra/test_astra_component.py +++ b/src/backend/tests/integration/astra/test_astra_component.py @@ -1,7 +1,5 @@ import os -from langchain_astradb import AstraDBVectorStore, CollectionVectorServiceOptions -from langflow.components.embeddings.AstraVectorize import AstraVectorizeComponent import pytest from integration.utils import MockEmbeddings, check_env_vars from langchain_core.documents import Document @@ -97,6 +95,10 @@ def test_astra_embeds_and_search(astra_fixture): reason="missing astra env vars", ) def test_astra_vectorize(): + from langchain_astradb import AstraDBVectorStore, CollectionVectorServiceOptions + + from langflow.components.embeddings.AstraVectorize import AstraVectorizeComponent + store = None try: options = {"provider": "nvidia", "modelName": "NV-Embed-QA", "parameters": {}, "authentication": {}} @@ -142,6 +144,10 @@ def test_astra_vectorize(): ) def test_astra_vectorize_with_provider_api_key(): """tests vectorize using an openai api key""" + from langchain_astradb import AstraDBVectorStore, CollectionVectorServiceOptions + + from langflow.components.embeddings.AstraVectorize import AstraVectorizeComponent + store = None try: application_token = os.getenv("ASTRA_DB_APPLICATION_TOKEN") @@ -186,6 +192,10 @@ def test_astra_vectorize_with_provider_api_key(): ) def test_astra_vectorize_passes_authentication(): """tests vectorize using the authentication parameter""" + from langchain_astradb import AstraDBVectorStore, CollectionVectorServiceOptions + + from langflow.components.embeddings.AstraVectorize import AstraVectorizeComponent + store = None try: application_token = os.getenv("ASTRA_DB_APPLICATION_TOKEN")