fix opensearch unit test

This constant was removed because it was causing a circular import and I missed removing the references to it in the unit tests.
apache · Apr 29, 2024 · d2fba72 · d2fba72
1 parent 68f762a
commit d2fba72
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 6 deletions.
diff --git a/airflow/providers/amazon/aws/triggers/opensearch_serverless.py b/airflow/providers/amazon/aws/triggers/opensearch_serverless.py
@@ -58,7 +58,7 @@ def __init__(
             status_message="Status of OpenSearch Serverless Collection is",
             status_queries=["status"],
             return_key="collection_id" if collection_id else "collection_name",
-            return_value=collection_id if collection_name else collection_name,
+            return_value=collection_id if collection_id else collection_name,
             waiter_delay=waiter_delay,
             waiter_max_attempts=waiter_max_attempts,
             aws_conn_id=aws_conn_id,

diff --git a/tests/providers/amazon/aws/sensors/test_opensearch_serverless.py b/tests/providers/amazon/aws/sensors/test_opensearch_serverless.py
@@ -79,7 +79,7 @@ def test_name_and_id_combinations(self, collection_name, collection_id, expected
             assert op.collection_name == collection_name
         if not expected_pass:
             with pytest.raises(
-                AttributeError, match=OpenSearchServerlessCollectionActiveSensor.INVALID_ARGS_MESSAGE
+                AttributeError, match="Either collection_ids or collection_names must be provided, not both."
             ):
                 OpenSearchServerlessCollectionActiveSensor(**call_args)
 

diff --git a/tests/providers/amazon/aws/triggers/test_opensearch_serverless.py b/tests/providers/amazon/aws/triggers/test_opensearch_serverless.py
@@ -22,9 +22,6 @@
 import pytest
 
 from airflow.providers.amazon.aws.hooks.opensearch_serverless import OpenSearchServerlessHook
-from airflow.providers.amazon.aws.sensors.opensearch_serverless import (
-    OpenSearchServerlessCollectionActiveSensor,
-)
 from airflow.providers.amazon.aws.triggers.opensearch_serverless import (
     OpenSearchServerlessCollectionActiveTrigger,
 )
@@ -72,7 +69,7 @@ def test_serialization(self, collection_name, collection_id, expected_pass):
 
         if not expected_pass:
             with pytest.raises(
-                AttributeError, match=OpenSearchServerlessCollectionActiveSensor.INVALID_ARGS_MESSAGE
+                AttributeError, match="Either collection_ids or collection_names must be provided, not both."
             ):
                 OpenSearchServerlessCollectionActiveTrigger(**call_args)