From e6750f449cbef07b37fe34fe76f1b4304b1421b6 Mon Sep 17 00:00:00 2001 From: vatsrahul1001 Date: Mon, 24 Jun 2024 11:40:10 +0530 Subject: [PATCH] add delete index teardown to pinecone system tests --- .../providers/pinecone/example_create_pod_index.py | 13 ++++++++++++- .../pinecone/example_create_serverless_index.py | 13 ++++++++++++- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/tests/system/providers/pinecone/example_create_pod_index.py b/tests/system/providers/pinecone/example_create_pod_index.py index 9b6f7d7d8882d..a2d7f16c6934e 100644 --- a/tests/system/providers/pinecone/example_create_pod_index.py +++ b/tests/system/providers/pinecone/example_create_pod_index.py @@ -20,6 +20,7 @@ from datetime import datetime from airflow import DAG +from airflow.decorators import task, teardown from airflow.providers.pinecone.operators.pinecone import CreatePodIndexOperator index_name = os.getenv("INDEX_NAME", "test") @@ -33,7 +34,7 @@ ) as dag: # [START howto_operator_create_pod_index] # reference: https://docs.pinecone.io/reference/api/control-plane/create_index - CreatePodIndexOperator( + create_index = CreatePodIndexOperator( task_id="pinecone_create_pod_index", index_name=index_name, dimension=3, @@ -44,6 +45,16 @@ ) # [END howto_operator_create_pod_index] + @teardown + @task + def delete_index(): + from airflow.providers.pinecone.hooks.pinecone import PineconeHook + + hook = PineconeHook() + hook.delete_index(index_name=index_name) + + create_index >> delete_index() + from tests.system.utils import get_test_run # noqa: E402 diff --git a/tests/system/providers/pinecone/example_create_serverless_index.py b/tests/system/providers/pinecone/example_create_serverless_index.py index a7924e63ef338..cf1e2c5cee4b5 100644 --- a/tests/system/providers/pinecone/example_create_serverless_index.py +++ b/tests/system/providers/pinecone/example_create_serverless_index.py @@ -20,6 +20,7 @@ from datetime import datetime from airflow import DAG +from airflow.decorators import task, teardown from airflow.providers.pinecone.operators.pinecone import CreateServerlessIndexOperator index_name = os.getenv("INDEX_NAME", "test") @@ -33,7 +34,7 @@ ) as dag: # [START howto_operator_create_serverless_index] # reference: https://docs.pinecone.io/reference/api/control-plane/create_index - CreateServerlessIndexOperator( + create_index = CreateServerlessIndexOperator( task_id="pinecone_create_serverless_index", index_name=index_name, dimension=128, @@ -43,6 +44,16 @@ ) # [END howto_operator_create_serverless_index] + @teardown + @task + def delete_index(): + from airflow.providers.pinecone.hooks.pinecone import PineconeHook + + hook = PineconeHook() + hook.delete_index(index_name=index_name) + + create_index >> delete_index() + from tests.system.utils import get_test_run # noqa: E402