This is an automated email from the ASF dual-hosted git repository.
ferruzzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f511c0d1246 Add a document_classifier_deletable waiter with unit tests
(#74125)
f511c0d1246 is described below
commit f511c0d124650cda5179c1e798d8afb70f553137
Author: D. Ferruzzi <[email protected]>
AuthorDate: Tue Oct 13 10:04:40 2026 -0700
Add a document_classifier_deletable waiter with unit tests (#74125)
* Add a document_classifier_deletable waiter with unit tests, and implement
it in the system test
---
.../providers/amazon/aws/waiters/comprehend.json | 55 ++++++++++++++++++++++
.../aws/example_comprehend_document_classifier.py | 21 +++++++--
.../unit/amazon/aws/waiters/test_comprehend.py | 54 +++++++++++++++++++++
3 files changed, 127 insertions(+), 3 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/waiters/comprehend.json
b/providers/amazon/src/airflow/providers/amazon/aws/waiters/comprehend.json
index 0d2d7259598..020cf189369 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/waiters/comprehend.json
+++ b/providers/amazon/src/airflow/providers/amazon/aws/waiters/comprehend.json
@@ -99,6 +99,61 @@
"state": "success"
}
]
+ },
+ "document_classifier_deletable": {
+ "delay": 10,
+ "maxAttempts": 30,
+ "operation": "DescribeDocumentClassifier",
+ "acceptors": [
+ {
+ "matcher": "path",
+ "argument": "DocumentClassifierProperties.Status",
+ "expected": "DELETING",
+ "state": "failure"
+ },
+ {
+ "matcher": "path",
+ "argument": "DocumentClassifierProperties.Status",
+ "expected": "SUBMITTED",
+ "state": "retry"
+ },
+ {
+ "matcher": "path",
+ "argument": "DocumentClassifierProperties.Status",
+ "expected": "TRAINING",
+ "state": "retry"
+ },
+ {
+ "matcher": "path",
+ "argument": "DocumentClassifierProperties.Status",
+ "expected": "STOP_REQUESTED",
+ "state": "retry"
+ },
+ {
+ "matcher": "path",
+ "argument": "DocumentClassifierProperties.Status",
+ "expected": "TRAINED",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "DocumentClassifierProperties.Status",
+ "expected": "TRAINED_WITH_WARNING",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "DocumentClassifierProperties.Status",
+ "expected": "IN_ERROR",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "DocumentClassifierProperties.Status",
+ "expected": "STOPPED",
+ "state": "success"
+ }
+ ]
}
}
}
diff --git
a/providers/amazon/tests/system/amazon/aws/example_comprehend_document_classifier.py
b/providers/amazon/tests/system/amazon/aws/example_comprehend_document_classifier.py
index 7b32eae9615..d15796165af 100644
---
a/providers/amazon/tests/system/amazon/aws/example_comprehend_document_classifier.py
+++
b/providers/amazon/tests/system/amazon/aws/example_comprehend_document_classifier.py
@@ -114,13 +114,28 @@ def document_classifier_workflow():
# [END howto_sensor_create_document_classifier]
@task(trigger_rule=TriggerRule.ALL_DONE)
- def delete_classifier(document_classifier_arn: str):
-
ComprehendHook().conn.delete_document_classifier(DocumentClassifierArn=document_classifier_arn)
+ def delete_classifier(classifier_name: str):
+ hook = ComprehendHook()
+ classifiers = hook.conn.list_document_classifiers(
+ Filter={"DocumentClassifierName": classifier_name},
+ )["DocumentClassifierPropertiesList"]
+
+ for classifier in classifiers:
+ classifier_arn = classifier["DocumentClassifierArn"]
+
+ # A classifier can only be deleted from a terminal state.
+ if classifier["Status"] == "DELETING":
+ continue
+ if classifier["Status"] in ("SUBMITTED", "TRAINING"):
+
hook.conn.stop_training_document_classifier(DocumentClassifierArn=classifier_arn)
+
+
hook.get_waiter("document_classifier_deletable").wait(DocumentClassifierArn=classifier_arn)
+
hook.conn.delete_document_classifier(DocumentClassifierArn=classifier_arn)
chain(
create_document_classifier,
await_create_document_classifier,
- delete_classifier(create_document_classifier.output),
+ delete_classifier(classifier_name),
)
diff --git a/providers/amazon/tests/unit/amazon/aws/waiters/test_comprehend.py
b/providers/amazon/tests/unit/amazon/aws/waiters/test_comprehend.py
index 650632a366c..05faecfc9b3 100644
--- a/providers/amazon/tests/unit/amazon/aws/waiters/test_comprehend.py
+++ b/providers/amazon/tests/unit/amazon/aws/waiters/test_comprehend.py
@@ -33,6 +33,7 @@ class TestComprehendCustomWaiters:
def test_service_waiters(self):
assert "pii_entities_detection_job_complete" in
ComprehendHook().list_waiters()
assert "create_document_classifier_complete" in
ComprehendHook().list_waiters()
+ assert "document_classifier_deletable" in
ComprehendHook().list_waiters()
class TestComprehendCustomWaitersBase:
@@ -102,3 +103,56 @@ class
TestComprehendDocumentClassifierCompleteWaiter(TestComprehendCustomWaiters
ComprehendHook().get_waiter(self.WAITER_NAME).wait(
DocumentClassifierArn="arn", WaiterConfig={"Delay": 0.01,
"MaxAttempts": 3}
)
+
+
+class
TestComprehendDocumentClassifierDeletableWaiter(TestComprehendCustomWaitersBase):
+ WAITER_NAME = "document_classifier_deletable"
+
+ @pytest.fixture
+ def mock_describe_document_classifier(self):
+ with mock.patch.object(self.client, "describe_document_classifier") as
mock_getter:
+ yield mock_getter
+
+ @pytest.mark.parametrize("state", ["TRAINED", "TRAINED_WITH_WARNING",
"IN_ERROR", "STOPPED"])
+ def test_document_classifier_deletable(self, state,
mock_describe_document_classifier):
+ mock_describe_document_classifier.return_value =
{"DocumentClassifierProperties": {"Status": state}}
+
+
ComprehendHook().get_waiter(self.WAITER_NAME).wait(DocumentClassifierArn="arn")
+
+
mock_describe_document_classifier.assert_called_once_with(DocumentClassifierArn="arn")
+
+ @pytest.mark.parametrize("state", ["SUBMITTED", "TRAINING",
"STOP_REQUESTED"])
+ def test_document_classifier_deletable_retries_while_in_progress(
+ self, state, mock_describe_document_classifier
+ ):
+ mock_describe_document_classifier.return_value =
{"DocumentClassifierProperties": {"Status": state}}
+
+ with pytest.raises(botocore.exceptions.WaiterError):
+ ComprehendHook().get_waiter(self.WAITER_NAME).wait(
+ DocumentClassifierArn="arn", WaiterConfig={"Delay": 0.01,
"MaxAttempts": 2}
+ )
+
+ assert mock_describe_document_classifier.call_count == 2
+
+ def test_document_classifier_deletable_fails_while_already_deleting(
+ self, mock_describe_document_classifier
+ ):
+ mock_describe_document_classifier.return_value = {
+ "DocumentClassifierProperties": {"Status": "DELETING"}
+ }
+
+ with pytest.raises(botocore.exceptions.WaiterError):
+ ComprehendHook().get_waiter(self.WAITER_NAME).wait(
+ DocumentClassifierArn="arn", WaiterConfig={"Delay": 0.01,
"MaxAttempts": 2}
+ )
+
+
mock_describe_document_classifier.assert_called_once_with(DocumentClassifierArn="arn")
+
+ def test_document_classifier_deletable_wait(self,
mock_describe_document_classifier):
+ wait = {"DocumentClassifierProperties": {"Status": "STOP_REQUESTED"}}
+ success = {"DocumentClassifierProperties": {"Status": "STOPPED"}}
+ mock_describe_document_classifier.side_effect = [wait, wait, success]
+
+ ComprehendHook().get_waiter(self.WAITER_NAME).wait(
+ DocumentClassifierArn="arn", WaiterConfig={"Delay": 0.01,
"MaxAttempts": 3}
+ )
