extract collection_name from system tests and make them unique

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()

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)
     )
 
 

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()

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,
             ]
         )
 

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

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()
 

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()