From 21f29bf69346f0569feb63fc76ec697138736459 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Tue, 2 Jul 2024 11:35:14 +0800 Subject: [PATCH] test(weaviate): extract collection_name from system tests and make them unique --- .../providers/weaviate/example_weaviate_cohere.py | 8 +++++--- .../weaviate/example_weaviate_dynamic_mapping_dag.py | 8 +++++--- .../providers/weaviate/example_weaviate_openai.py | 12 ++++++------ .../providers/weaviate/example_weaviate_operator.py | 10 ++++++---- .../weaviate/example_weaviate_using_hook.py | 8 +++++--- .../weaviate/example_weaviate_vectorizer_dag.py | 8 ++++---- .../example_weaviate_without_vectorizer_dag.py | 8 ++++---- 7 files changed, 35 insertions(+), 27 deletions(-) diff --git a/tests/system/providers/weaviate/example_weaviate_cohere.py b/tests/system/providers/weaviate/example_weaviate_cohere.py index 8415bcecc7ac8..d8662e464bd4b 100644 --- a/tests/system/providers/weaviate/example_weaviate_cohere.py +++ b/tests/system/providers/weaviate/example_weaviate_cohere.py @@ -22,6 +22,8 @@ from airflow.providers.cohere.operators.embedding import CohereEmbeddingOperator from airflow.providers.weaviate.operators.weaviate import WeaviateIngestOperator +COLLECTION_NAME = "weaviate_cohere_example_collection" + @dag( schedule=None, @@ -44,7 +46,7 @@ def create_weaviate_collection(): weaviate_hook = WeaviateHook() # Collection definition object. Weaviate's autoschema feature will infer properties when importing. - weaviate_hook.create_collection(name="Weaviate_example_collection", vectorizer_config=None) + weaviate_hook.create_collection(name=COLLECTION_NAME, vectorizer_config=None) @setup @task @@ -78,7 +80,7 @@ def update_vector_data_in_json(**kwargs): perform_ingestion = WeaviateIngestOperator( task_id="perform_ingestion", conn_id="weaviate_default", - collection_name="Weaviate_example_collection", + collection_name=COLLECTION_NAME, input_data=update_vector_data_in_json["return_value"], ) @@ -98,7 +100,7 @@ def delete_weaviate_collections(): weaviate_hook = WeaviateHook() # collection definition object. Weaviate's autoschema feature will infer properties when importing. - weaviate_hook.delete_collections(["Weaviate_example_collections"]) + weaviate_hook.delete_collections([COLLECTION_NAME]) ( create_weaviate_collection() diff --git a/tests/system/providers/weaviate/example_weaviate_dynamic_mapping_dag.py b/tests/system/providers/weaviate/example_weaviate_dynamic_mapping_dag.py index a603287863904..5f998cc52e862 100644 --- a/tests/system/providers/weaviate/example_weaviate_dynamic_mapping_dag.py +++ b/tests/system/providers/weaviate/example_weaviate_dynamic_mapping_dag.py @@ -22,6 +22,8 @@ from airflow.decorators import dag, setup, task, teardown from airflow.providers.weaviate.operators.weaviate import WeaviateIngestOperator +COLLECTION_NAMES = ["Weaviate_DTM_example_collection_1", "Weaviate_DTM_example_collection_2"] + @dag( schedule=None, @@ -61,7 +63,7 @@ def get_data_to_ingest(): task_id="perform_ingestion", conn_id="weaviate_default", ).expand( - collection_name=["example1", "example2"], + collection_name=COLLECTION_NAMES, input_data=get_data_to_ingest["return_value"], ) @@ -81,10 +83,10 @@ def delete_weaviate_collection(collection_name): ( create_weaviate_collection.expand( - data=[["example1", None], ["example2", Configure.Vectorizer.text2vec_openai()]] + data=[[COLLECTION_NAMES[0], None], [COLLECTION_NAMES[1], Configure.Vectorizer.text2vec_openai()]] ) >> perform_ingestion - >> delete_weaviate_collection.expand(collection_name=["example1", "example2"]) + >> delete_weaviate_collection.expand(collection_name=COLLECTION_NAMES) ) diff --git a/tests/system/providers/weaviate/example_weaviate_openai.py b/tests/system/providers/weaviate/example_weaviate_openai.py index 31fe0e8e29ab9..adf20fd929f1b 100644 --- a/tests/system/providers/weaviate/example_weaviate_openai.py +++ b/tests/system/providers/weaviate/example_weaviate_openai.py @@ -26,6 +26,8 @@ from airflow.providers.weaviate.hooks.weaviate import WeaviateHook from airflow.providers.weaviate.operators.weaviate import WeaviateIngestOperator +COLLECTION_NAME = "Weaviate_openai_example_collection" + @dag( schedule=None, @@ -46,7 +48,7 @@ def create_weaviate_collection(): """ weaviate_hook = WeaviateHook() # collection definition object. Weaviate's autoschema feature will infer properties when importing. - weaviate_hook.create_collection("Weaviate_openai_example_collection") + weaviate_hook.create_collection(COLLECTION_NAME) @setup @task @@ -76,7 +78,7 @@ def update_vector_data_in_json(**kwargs): perform_ingestion = WeaviateIngestOperator( task_id="perform_ingestion", conn_id="weaviate_default", - collection_name="Weaviate_openai_example_collection", + collection_name=COLLECTION_NAME, input_data=update_vector_data_in_json["return_value"], ) @@ -93,9 +95,7 @@ def query_weaviate(**kwargs): query_vector = ti.xcom_pull(task_ids="embed_query", key="return_value") weaviate_hook = WeaviateHook() properties = ["question", "answer", "category"] - response = weaviate_hook.query_with_vector( - query_vector, "Weaviate_openai_example_collection", properties - ) + response = weaviate_hook.query_with_vector(query_vector, COLLECTION_NAME, properties) assert "In 1953 Watson & Crick built a model" in response.objects[0].properties["question"] @teardown @@ -107,7 +107,7 @@ def delete_weaviate_collection(): weaviate_hook = WeaviateHook() # collection definition object. Weaviate's autoschema feature will infer properties when importing. - weaviate_hook.delete_collections(["Weaviate_openai_example_collection"]) + weaviate_hook.delete_collections([COLLECTION_NAME]) ( create_weaviate_collection() diff --git a/tests/system/providers/weaviate/example_weaviate_operator.py b/tests/system/providers/weaviate/example_weaviate_operator.py index b92538059e556..081f2ef3903d7 100644 --- a/tests/system/providers/weaviate/example_weaviate_operator.py +++ b/tests/system/providers/weaviate/example_weaviate_operator.py @@ -26,6 +26,8 @@ WeaviateIngestOperator, ) +COLLECTION_NAME = "QuestionWithoutVectorizerUsingOperator" + sample_data_with_vector = [ { "Answer": "Liver", @@ -108,7 +110,7 @@ def create_collection_without_vectorizer(): weaviate_hook = WeaviateHook() # collection definition object. Weaviate's autoschema feature will infer properties when importing. - weaviate_hook.create_collection("QuestionWithoutVectorizerUsingOperator") + weaviate_hook.create_collection(COLLECTION_NAME) @task(trigger_rule="all_done") def store_data_with_vectors_in_xcom(): @@ -118,7 +120,7 @@ def store_data_with_vectors_in_xcom(): batch_data_with_vectors_xcom_data = WeaviateIngestOperator( task_id="batch_data_with_vectors_xcom_data", conn_id="weaviate_default", - collection_name="QuestionWithoutVectorizerUsingOperator", + collection_name=COLLECTION_NAME, input_data=store_data_with_vectors_in_xcom(), trigger_rule="all_done", ) @@ -128,7 +130,7 @@ def store_data_with_vectors_in_xcom(): batch_data_with_vectors_callable_data = WeaviateIngestOperator( task_id="batch_data_with_vectors_callable_data", conn_id="weaviate_default", - collection_name="QuestionWithoutVectorizerUsingOperator", + collection_name=COLLECTION_NAME, input_data=get_data_with_vectors(), trigger_rule="all_done", ) @@ -255,7 +257,7 @@ def delete_weaviate_collection_without_vector(): weaviate_hook.delete_collections( [ - "QuestionWithoutVectorizerUsingOperator", + COLLECTION_NAME, ] ) diff --git a/tests/system/providers/weaviate/example_weaviate_using_hook.py b/tests/system/providers/weaviate/example_weaviate_using_hook.py index 11d57bfbef2a9..8ad7356465a48 100644 --- a/tests/system/providers/weaviate/example_weaviate_using_hook.py +++ b/tests/system/providers/weaviate/example_weaviate_using_hook.py @@ -22,6 +22,8 @@ from airflow.decorators import dag, task, teardown +COLLECTION_NAME = "QuestionWithOpenAIVectorizerUsingHook" + @dag( schedule=None, @@ -41,7 +43,7 @@ def create_collection_with_vectorizer(): weaviate_hook = WeaviateHook() weaviate_hook.create_collection( - "QuestionWithOpenAIVectorizerUsingHook", + COLLECTION_NAME, description="Information from a Jeopardy! question", properties=[ Property(name="question", description="The question", data_type=DataType.TEXT), @@ -86,7 +88,7 @@ def batch_data_without_vectors(data: list): from airflow.providers.weaviate.hooks.weaviate import WeaviateHook weaviate_hook = WeaviateHook() - weaviate_hook.batch_data("QuestionWithOpenAIVectorizerUsingHook", data) + weaviate_hook.batch_data(COLLECTION_NAME, data) @task(trigger_rule="all_done") def batch_data_with_vectors(data: list): @@ -106,7 +108,7 @@ def delete_weaviate_collection_vector(): weaviate_hook = WeaviateHook() # collection definition object. Weaviate's autoschema feature will infer properties when importing. - weaviate_hook.delete_collections(["QuestionWithOpenAIVectorizerUsingHook"]) + weaviate_hook.delete_collections([COLLECTION_NAME]) @teardown @task diff --git a/tests/system/providers/weaviate/example_weaviate_vectorizer_dag.py b/tests/system/providers/weaviate/example_weaviate_vectorizer_dag.py index abd3a3de85193..546148ad70ed2 100644 --- a/tests/system/providers/weaviate/example_weaviate_vectorizer_dag.py +++ b/tests/system/providers/weaviate/example_weaviate_vectorizer_dag.py @@ -22,7 +22,7 @@ from airflow.decorators import dag, setup, task, teardown from airflow.providers.weaviate.operators.weaviate import WeaviateIngestOperator -collection_name = "Weaviate_with_vectorizer_example_collection" +COLLECTION_NAME = "Weaviate_with_vectorizer_example_collection" @dag( @@ -47,7 +47,7 @@ def create_weaviate_collection(): weaviate_hook = WeaviateHook() # collection definition object. Weaviate's autoschema feature will infer properties when importing. weaviate_hook.create_collection( - collection_name, + COLLECTION_NAME, vectorizer_config=Configure.Vectorizer.text2vec_openai(), ) @@ -65,7 +65,7 @@ def get_data_to_ingest(): perform_ingestion = WeaviateIngestOperator( task_id="perform_ingestion", conn_id="weaviate_default", - collection_name=collection_name, + collection_name=COLLECTION_NAME, input_data=data_to_ingest["return_value"], ) @@ -91,7 +91,7 @@ def delete_weaviate_collection(): weaviate_hook = WeaviateHook() # collection definition object. Weaviate's autoschema feature will infer properties when importing. - weaviate_hook.delete_collections([collection_name]) + weaviate_hook.delete_collections([COLLECTION_NAME]) create_weaviate_collection() >> perform_ingestion >> query_weaviate() >> delete_weaviate_collection() diff --git a/tests/system/providers/weaviate/example_weaviate_without_vectorizer_dag.py b/tests/system/providers/weaviate/example_weaviate_without_vectorizer_dag.py index 5c97020bbad86..c4062508e26d1 100644 --- a/tests/system/providers/weaviate/example_weaviate_without_vectorizer_dag.py +++ b/tests/system/providers/weaviate/example_weaviate_without_vectorizer_dag.py @@ -22,7 +22,7 @@ from airflow.providers.openai.operators.openai import OpenAIEmbeddingOperator from airflow.providers.weaviate.operators.weaviate import WeaviateIngestOperator -collection_name = "Weaviate_example_without_vectorizer_collection" +COLLECTION_NAME = "Weaviate_example_without_vectorizer_collection" @dag( @@ -46,7 +46,7 @@ def create_weaviate_collection(): weaviate_hook = WeaviateHook() # collection definition object. Weaviate's autoschema feature will infer properties when importing. - weaviate_hook.create_collection(collection_name, vectorizer_config=None) + weaviate_hook.create_collection(COLLECTION_NAME, vectorizer_config=None) @setup @task @@ -62,7 +62,7 @@ def get_data_without_vectors(): perform_ingestion = WeaviateIngestOperator( task_id="perform_ingestion", conn_id="weaviate_default", - collection_name=collection_name, + collection_name=COLLECTION_NAME, input_data=data_to_ingest["return_value"], ) @@ -97,7 +97,7 @@ def delete_weaviate_collection(): weaviate_hook = WeaviateHook() # collection definition object. Weaviate's autoschema feature will infer properties when importing. - weaviate_hook.delete_collections([collection_name]) + weaviate_hook.delete_collections([COLLECTION_NAME]) ( create_weaviate_collection()