cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r3067627555
##########
providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py:
##########
@@ -642,6 +644,317 @@ def test_execute_with_identity_dict(self, aci_mock):
# user_assigned_identities should contain the resource id as a key
assert resource_id in (called_cg.identity.user_assigned_identities or
{})
+ @mock.patch(
+
"airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_deferrable_defers_when_container_running(self, aci_mock):
+ """When deferrable=True and the container is still running, defer() is
called."""
+ running_cg = make_mock_container(state="Running", exit_code=0,
detail_status="test")
+ aci_mock.return_value.get_state.return_value = running_cg
+ aci_mock.return_value.exists.return_value = False
+
+ aci = AzureContainerInstancesOperator(
+ ci_conn_id="azure_default",
+ registry_conn_id=None,
+ resource_group="resource-group",
+ name="container-name",
+ image="container-image",
+ region="region",
+ task_id="task",
+ deferrable=True,
+ polling_interval=10.0,
+ )
+ with pytest.raises(TaskDeferred) as exc_info:
+ aci.execute(None)
+
+ assert isinstance(exc_info.value.trigger,
AzureContainerInstanceTrigger)
+ assert exc_info.value.trigger.resource_group == "resource-group"
+ assert exc_info.value.trigger.name == "container-name"
+ assert exc_info.value.trigger.polling_interval == 10.0
+ assert exc_info.value.method_name == "execute_complete"
+ # Container must NOT be deleted when deferring — it is still running
+ assert aci_mock.return_value.delete.call_count == 0
+
+ @mock.patch(
+
"airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def
test_execute_deferrable_completes_synchronously_if_already_terminated(self,
aci_mock):
+ """When deferrable=True but container is already terminal, no deferral
— sync completion."""
+ terminated_cg = make_mock_container(state="Terminated", exit_code=0,
detail_status="test")
+ aci_mock.return_value.get_state.return_value = terminated_cg
+ aci_mock.return_value.exists.return_value = False
+
+ aci = AzureContainerInstancesOperator(
+ ci_conn_id="azure_default",
+ registry_conn_id=None,
+ resource_group="resource-group",
+ name="container-name",
+ image="container-image",
+ region="region",
+ task_id="task",
+ deferrable=True,
+ )
+ result = aci.execute(None)
+ assert result == 0
+ assert aci_mock.return_value.create_or_update.call_count == 1
+
+ @mock.patch(
+
"airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_complete_success(self, aci_mock):
+ """execute_complete succeeds, does not raise, and deletes the
container group."""
+ aci_mock.return_value.get_logs.return_value = None
+
+ aci = AzureContainerInstancesOperator(
+ ci_conn_id="azure_default",
+ registry_conn_id=None,
+ resource_group="resource-group",
+ name="container-name",
+ image="container-image",
+ region="region",
+ task_id="task",
+ deferrable=True,
+ remove_on_success=True,
+ )
+ result = aci.execute_complete(
+ context=None,
+ event={
+ "status": "success",
+ "exit_code": 0,
+ "detail_status": "Completed",
+ "resource_group": "resource-group",
+ "name": "container-name",
+ },
+ )
+ assert result == 0
+ assert aci_mock.return_value.delete.call_count == 1
+
+ @mock.patch(
+
"airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_complete_success_with_remove_on_success_false(self,
aci_mock):
+ """execute_complete with remove_on_success=False should NOT delete the
container."""
+ aci_mock.return_value.get_logs.return_value = None
+
+ aci = AzureContainerInstancesOperator(
+ ci_conn_id="azure_default",
+ registry_conn_id=None,
+ resource_group="resource-group",
+ name="container-name",
+ image="container-image",
+ region="region",
+ task_id="task",
+ deferrable=True,
+ remove_on_success=False,
+ )
+ aci.execute_complete(
+ context=None,
+ event={
+ "status": "success",
+ "exit_code": 0,
+ "detail_status": "Completed",
+ "resource_group": "resource-group",
+ "name": "container-name",
+ },
+ )
+ assert aci_mock.return_value.delete.call_count == 0
+
+ @mock.patch(
+
"airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_complete_error_event_raises(self, aci_mock):
+ """execute_complete raises AirflowException when event has
status=error."""
Review Comment:
Stale docstring, will correct it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
