This is an automated email from the ASF dual-hosted git repository.
potiuk 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 dcb19a9103a Add deferrable mode to AzureContainerInstancesOperator
(#62772)
dcb19a9103a is described below
commit dcb19a9103aff5f0c2315538ed9eb3374a471561
Author: Akshaykumar <[email protected]>
AuthorDate: Sun May 10 07:34:43 2026 +0530
Add deferrable mode to AzureContainerInstancesOperator (#62772)
* Add aync support to ACI
* avoid raising AirflowException directly
* Address review
- async credential: store and close _async_credential in close()
- Remove per-poll state logging in trigger
- Change default polling_interval from 5s to 30s
- Simplify finally cleanup block in execute()
- Add get_async_conn() caching test
* address review comments
* Refactor async hook to use asynccontextmanager
* Fix CI
* ensure proper execution of cleanup
* address review
---------
Co-authored-by: Akshay <[email protected]>
---
providers/microsoft/azure/provider.yaml | 3 +
.../providers/microsoft/azure/get_provider_info.py | 4 +
.../microsoft/azure/hooks/container_instance.py | 96 ++++++-
.../azure/operators/container_instances.py | 92 ++++++-
.../microsoft/azure/triggers/container_instance.py | 130 ++++++++++
.../azure/hooks/test_container_instance.py | 265 ++++++++++++++++++-
.../azure/operators/test_container_instances.py | 288 +++++++++++++++++++++
.../azure/triggers/test_container_instance.py | 248 ++++++++++++++++++
8 files changed, 1111 insertions(+), 15 deletions(-)
diff --git a/providers/microsoft/azure/provider.yaml
b/providers/microsoft/azure/provider.yaml
index 72ab88da6fe..76421ba3127 100644
--- a/providers/microsoft/azure/provider.yaml
+++ b/providers/microsoft/azure/provider.yaml
@@ -312,6 +312,9 @@ triggers:
- integration-name: Microsoft Azure Compute
python-modules:
- airflow.providers.microsoft.azure.triggers.compute
+ - integration-name: Microsoft Azure Container Instances
+ python-modules:
+ - airflow.providers.microsoft.azure.triggers.container_instance
- integration-name: Microsoft Azure Data Factory
python-modules:
- airflow.providers.microsoft.azure.triggers.data_factory
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
index 944b296d708..0b312ec09a9 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py
@@ -288,6 +288,10 @@ def get_provider_info():
"integration-name": "Microsoft Azure Compute",
"python-modules":
["airflow.providers.microsoft.azure.triggers.compute"],
},
+ {
+ "integration-name": "Microsoft Azure Container Instances",
+ "python-modules":
["airflow.providers.microsoft.azure.triggers.container_instance"],
+ },
{
"integration-name": "Microsoft Azure Data Factory",
"python-modules":
["airflow.providers.microsoft.azure.triggers.data_factory"],
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py
index 27941ec53af..bf7e7177298 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py
@@ -17,17 +17,30 @@
# under the License.
from __future__ import annotations
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
from functools import cached_property
from typing import TYPE_CHECKING, Any, cast
from azure.common.client_factory import get_client_from_auth_file,
get_client_from_json_dict
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import ClientSecretCredential, DefaultAzureCredential
+from azure.identity.aio import (
+ ClientSecretCredential as AsyncClientSecretCredential,
+ DefaultAzureCredential as AsyncDefaultAzureCredential,
+)
from azure.mgmt.containerinstance import ContainerInstanceManagementClient
+from azure.mgmt.containerinstance.aio import (
+ ContainerInstanceManagementClient as
AsyncContainerInstanceManagementClient,
+)
+from airflow.providers.common.compat.connection import get_async_connection
from airflow.providers.common.compat.sdk import AirflowException
from airflow.providers.microsoft.azure.hooks.base_azure import
_AZURE_CLOUD_ENVIRONMENTS, AzureBaseHook
-from airflow.providers.microsoft.azure.utils import
get_sync_default_azure_credential
+from airflow.providers.microsoft.azure.utils import (
+ get_async_default_azure_credential,
+ get_sync_default_azure_credential,
+)
if TYPE_CHECKING:
from azure.mgmt.containerinstance.models import (
@@ -180,3 +193,84 @@ class AzureContainerInstanceHook(AzureBaseHook):
return False, str(e)
return True, "Successfully connected to Azure Container Instance."
+
+
+class AzureContainerInstanceAsyncHook(AzureContainerInstanceHook):
+ """
+ An async hook for communicating with Azure Container Instances.
+
+ :param azure_conn_id: :ref:`Azure connection id<howto/connection:azure>` of
+ a service principal which will be used to start the container instance.
+ """
+
+ def __init__(self, azure_conn_id: str =
AzureContainerInstanceHook.default_conn_name) -> None:
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ @asynccontextmanager
+ async def get_async_conn(self) ->
AsyncGenerator[AsyncContainerInstanceManagementClient, None]:
+ """Create an async management client bound to a single credential."""
+ conn = await get_async_connection(self.conn_id)
+ tenant = conn.extra_dejson.get("tenantId")
+ subscription_id = cast("str", conn.extra_dejson.get("subscriptionId"))
+
+ credential: AsyncClientSecretCredential | AsyncDefaultAzureCredential
+ if all([conn.login, conn.password, tenant]):
+ credential = AsyncClientSecretCredential(
+ client_id=cast("str", conn.login),
+ client_secret=cast("str", conn.password),
+ tenant_id=cast("str", tenant),
+ )
+ else:
+ managed_identity_client_id =
conn.extra_dejson.get("managed_identity_client_id")
+ workload_identity_tenant_id =
conn.extra_dejson.get("workload_identity_tenant_id")
+ credential = get_async_default_azure_credential(
+ managed_identity_client_id=managed_identity_client_id,
+ workload_identity_tenant_id=workload_identity_tenant_id,
+ )
+
+ client = AsyncContainerInstanceManagementClient(
+ credential=credential,
+ subscription_id=subscription_id,
+ )
+ try:
+ yield client
+ finally:
+ await client.close()
+ if hasattr(credential, "close"):
+ await credential.close()
+
+ async def get_state(self, resource_group: str, name: str) ->
ContainerGroup: # type: ignore[override]
+ """
+ Get the state of a container group asynchronously.
+
+ :param resource_group: the name of the resource group
+ :param name: the name of the container group
+ :return: ContainerGroup
+ """
+ async with self.get_async_conn() as client:
+ return await client.container_groups.get(resource_group, name)
+
+ async def get_logs(self, resource_group: str, name: str, tail: int = 1000)
-> list: # type: ignore[override]
+ """
+ Get the tail from logs of a container group asynchronously.
+
+ :param resource_group: the name of the resource group
+ :param name: the name of the container group
+ :param tail: the size of the tail
+ :return: A list of log messages
+ """
+ async with self.get_async_conn() as client:
+ logs = await client.containers.list_logs(resource_group, name,
name, tail=tail)
+ if logs.content is None:
+ return [None]
+ return logs.content.splitlines(True)
+
+ async def delete(self, resource_group: str, name: str) -> None: # type:
ignore[override]
+ """
+ Delete a container group asynchronously.
+
+ :param resource_group: the name of the resource group
+ :param name: the name of the container group
+ """
+ async with self.get_async_conn() as client:
+ await client.container_groups.begin_delete(resource_group, name)
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py
index 3bcacfd4c00..3901ff5258e 100644
---
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py
@@ -21,6 +21,7 @@ import re
import time
from collections import namedtuple
from collections.abc import Sequence
+from functools import cached_property
from typing import TYPE_CHECKING, Any, cast
from azure.mgmt.containerinstance.models import (
@@ -42,10 +43,11 @@ from azure.mgmt.containerinstance.models import (
)
from msrestazure.azure_exceptions import CloudError
-from airflow.providers.common.compat.sdk import AirflowException,
AirflowTaskTimeout, BaseOperator
+from airflow.providers.common.compat.sdk import AirflowException,
AirflowTaskTimeout, BaseOperator, conf
from airflow.providers.microsoft.azure.hooks.container_instance import
AzureContainerInstanceHook
from airflow.providers.microsoft.azure.hooks.container_registry import
AzureContainerRegistryHook
from airflow.providers.microsoft.azure.hooks.container_volume import
AzureContainerVolumeHook
+from airflow.providers.microsoft.azure.triggers.container_instance import
AzureContainerInstanceTrigger
if TYPE_CHECKING:
from airflow.sdk import Context
@@ -105,6 +107,10 @@ class AzureContainerInstancesOperator(BaseOperator):
:param diagnostics: Container group diagnostic information (Log Analytics).
:param priority: Container group priority, Possible values include:
'Regular', 'Spot'
:param identity: List of User/System assigned identities for the container
group.
+ :param deferrable: Run in deferrable mode, releasing the worker slot while
the container
+ runs. Defaults to ``[operators] default_deferrable`` in
``airflow.cfg``.
+ :param remove_on_success: Delete the container group after a successful
run. Default ``True``.
+ :param polling_interval: Seconds between status polls in deferrable mode.
Default ``30.0``.
**Example**::
@@ -181,6 +187,7 @@ class AzureContainerInstancesOperator(BaseOperator):
gpu: Any | None = None,
command: list[str] | None = None,
remove_on_error: bool = True,
+ remove_on_success: bool = True,
fail_if_exists: bool = True,
tags: dict[str, str] | None = None,
xcom_all: bool | None = None,
@@ -193,6 +200,8 @@ class AzureContainerInstancesOperator(BaseOperator):
diagnostics: ContainerGroupDiagnostics | None = None,
priority: str | None = "Regular",
identity: ContainerGroupIdentity | dict | None = None,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ polling_interval: float = 30.0,
**kwargs,
) -> None:
super().__init__(**kwargs)
@@ -211,8 +220,8 @@ class AzureContainerInstancesOperator(BaseOperator):
self.gpu = gpu
self.command = command
self.remove_on_error = remove_on_error
+ self.remove_on_success = remove_on_success
self.fail_if_exists = fail_if_exists
- self._ci_hook: Any = None
self.tags = tags
self.xcom_all = xcom_all
self.os_type = os_type
@@ -242,6 +251,8 @@ class AzureContainerInstancesOperator(BaseOperator):
"Please set 'Regular' or 'Spot' as the priority. "
f"Found `{self.priority}`."
)
+ self.deferrable = deferrable
+ self.polling_interval = polling_interval
# helper to accept dict (user-friendly) or ContainerGroupIdentity (SDK
object)
@staticmethod
@@ -303,15 +314,17 @@ class AzureContainerInstancesOperator(BaseOperator):
)
return identity
+ @cached_property
+ def hook(self) -> AzureContainerInstanceHook:
+ return AzureContainerInstanceHook(azure_conn_id=self.ci_conn_id)
+
def execute(self, context: Context) -> int:
# Check name again in case it was templated.
self._check_name(self.name)
- self._ci_hook =
AzureContainerInstanceHook(azure_conn_id=self.ci_conn_id)
-
if self.fail_if_exists:
self.log.info("Testing if container group already exists")
- if self._ci_hook.exists(self.resource_group, self.name):
+ if self.hook.exists(self.resource_group, self.name):
raise AirflowException("Container group exists")
if self.registry_conn_id:
@@ -340,6 +353,7 @@ class AzureContainerInstancesOperator(BaseOperator):
volume_mounts.append(VolumeMount(name=mount_name,
mount_path=mount_path, read_only=read_only))
exit_code = 1
+ _cleanup = True
try:
self.log.info("Starting container group with %.1f cpu %.1f mem",
self.cpu, self.memory_in_gb)
if self.gpu:
@@ -381,13 +395,26 @@ class AzureContainerInstancesOperator(BaseOperator):
identity=self.identity,
)
- self._ci_hook.create_or_update(self.resource_group, self.name,
container_group)
+ self.hook.create_or_update(self.resource_group, self.name,
container_group)
self.log.info("Container group started %s/%s",
self.resource_group, self.name)
+ if self.deferrable:
+ _cleanup = False
+ self.defer(
+ trigger=AzureContainerInstanceTrigger(
+ resource_group=self.resource_group,
+ name=self.name,
+ ci_conn_id=self.ci_conn_id,
+ polling_interval=self.polling_interval,
+ ),
+ method_name=self.execute_complete.__name__,
+ timeout=self.execution_timeout,
+ )
+
exit_code = self._monitor_logging(self.resource_group, self.name)
if self.xcom_all is not None:
- logs = self._ci_hook.get_logs(self.resource_group, self.name)
+ logs = self.hook.get_logs(self.resource_group, self.name)
if logs is None:
context["ti"].xcom_push(key="logs", value=[])
else:
@@ -407,16 +434,57 @@ class AzureContainerInstancesOperator(BaseOperator):
raise AirflowException("Could not start container group")
finally:
- if exit_code == 0 or self.remove_on_error:
- self.on_kill()
+ if _cleanup:
+ if exit_code == 0 and self.remove_on_success:
+ self.on_kill()
+ elif exit_code != 0 and self.remove_on_error:
+ self.on_kill()
def on_kill(self) -> None:
self.log.info("Deleting container group")
try:
- self._ci_hook.delete(self.resource_group, self.name)
+ self.hook.delete(self.resource_group, self.name)
except Exception:
self.log.exception("Could not delete container group")
+ def execute_complete(self, context: Context, event: dict[str, Any] | None)
-> int:
+ """
+ Handle the trigger event after deferral.
+
+ Called by the Triggerer when the container reaches a terminal state.
+ Raises on failure; returns the exit code on success.
+ """
+ if event is None:
+ raise ValueError("Trigger error: event is None")
+
+ exit_code: int = event.get("exit_code", 1)
+
+ if event["status"] == "error":
+ if self.remove_on_error:
+ self.on_kill()
+ raise RuntimeError(
+ event.get(
+ "message",
+ f"Container group {self.resource_group}/{self.name} failed
with exit code {exit_code}",
+ )
+ )
+
+ try:
+ if self.xcom_all is not None:
+ logs = self.hook.get_logs(self.resource_group, self.name)
+ if logs is None:
+ context["ti"].xcom_push(key="logs", value=[])
+ elif self.xcom_all:
+ context["ti"].xcom_push(key="logs", value=logs)
+ else:
+ context["ti"].xcom_push(key="logs", value=logs[-1:])
+
+ self.log.info("Container had exit code: %s", exit_code)
+ return exit_code
+ finally:
+ if self.remove_on_success:
+ self.on_kill()
+
def _monitor_logging(self, resource_group: str, name: str) -> int:
last_state = None
last_message_logged = None
@@ -424,7 +492,7 @@ class AzureContainerInstancesOperator(BaseOperator):
while True:
try:
- cg_state = self._ci_hook.get_state(resource_group, name)
+ cg_state = self.hook.get_state(resource_group, name)
instance_view = cg_state.containers[0].instance_view
# If there is no instance view, we show the provisioning state
if instance_view is not None:
@@ -449,7 +517,7 @@ class AzureContainerInstancesOperator(BaseOperator):
if state in ["Running", "Terminated", "Succeeded"]:
try:
- logs = self._ci_hook.get_logs(resource_group, name)
+ logs = self.hook.get_logs(resource_group, name)
if logs and logs[0] is None:
self.log.error("Container log is broken, marking
as failed.")
return 1
diff --git
a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py
new file mode 100644
index 00000000000..db855d1b0e3
--- /dev/null
+++
b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py
@@ -0,0 +1,130 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncIterator
+from typing import Any
+
+from airflow.providers.microsoft.azure.hooks.container_instance import
AzureContainerInstanceAsyncHook
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+TERMINAL_STATES = frozenset({"Terminated", "Succeeded", "Failed", "Unhealthy"})
+SUCCESS_STATES = frozenset({"Terminated", "Succeeded"})
+
+
+class AzureContainerInstanceTrigger(BaseTrigger):
+ """
+ Poll an Azure Container Instance until it reaches a terminal state.
+
+ :param resource_group: the name of the resource group
+ :param name: the name of the container group
+ :param ci_conn_id: connection id of the Azure service principal
+ :param polling_interval: time in seconds between state polls
+ """
+
+ def __init__(
+ self,
+ resource_group: str,
+ name: str,
+ ci_conn_id: str,
+ polling_interval: float = 30.0,
+ ) -> None:
+ super().__init__()
+ self.resource_group = resource_group
+ self.name = name
+ self.ci_conn_id = ci_conn_id
+ self.polling_interval = polling_interval
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ """Serialize trigger arguments and classpath."""
+ return (
+
"airflow.providers.microsoft.azure.triggers.container_instance.AzureContainerInstanceTrigger",
+ {
+ "resource_group": self.resource_group,
+ "name": self.name,
+ "ci_conn_id": self.ci_conn_id,
+ "polling_interval": self.polling_interval,
+ },
+ )
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ """Poll ACI until a terminal state is reached, then yield a
TriggerEvent."""
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=self.ci_conn_id)
+ try:
+ async with hook.get_async_conn() as client:
+ while True:
+ cg_state = await
client.container_groups.get(self.resource_group, self.name)
+ instance_view = cg_state.containers[0].instance_view
+
+ if instance_view is not None:
+ c_state = instance_view.current_state
+ state = c_state.state
+ exit_code = c_state.exit_code
+ detail_status = c_state.detail_status
+ else:
+ prov = cg_state.provisioning_state
+ if prov in ("Failed", "Unhealthy"):
+ state = prov
+ exit_code = 1
+ detail_status = prov
+ else:
+ await asyncio.sleep(self.polling_interval)
+ continue
+
+ if state in TERMINAL_STATES:
+ if state in SUCCESS_STATES and exit_code == 0:
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "exit_code": exit_code,
+ "detail_status": detail_status,
+ "resource_group": self.resource_group,
+ "name": self.name,
+ }
+ )
+ else:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "exit_code": exit_code,
+ "detail_status": detail_status,
+ "resource_group": self.resource_group,
+ "name": self.name,
+ "message": (
+ f"Container group
{self.resource_group}/{self.name} "
+ f"reached state {state!r} with exit
code {exit_code} "
+ f"({detail_status})"
+ ),
+ }
+ )
+ return
+
+ await asyncio.sleep(self.polling_interval)
+ except asyncio.CancelledError:
+ raise
+ except Exception as e:
+ yield TriggerEvent(
+ {
+ "status": "error",
+ "message": str(e),
+ "resource_group": self.resource_group,
+ "name": self.name,
+ "exit_code": 1,
+ "detail_status": "",
+ }
+ )
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py
index 0bebf4d7afc..33df9b1d992 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py
@@ -17,19 +17,27 @@
# under the License.
from __future__ import annotations
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import AzureAuthorityHosts
+from azure.identity.aio import ClientSecretCredential as
AsyncClientSecretCredential
+from azure.mgmt.containerinstance.aio import (
+ ContainerInstanceManagementClient as
AsyncContainerInstanceManagementClient,
+)
from azure.mgmt.containerinstance.models import (
+ ContainerGroup,
Logs,
ResourceRequests,
ResourceRequirements,
)
from airflow.models import Connection
-from airflow.providers.microsoft.azure.hooks.container_instance import
AzureContainerInstanceHook
+from airflow.providers.microsoft.azure.hooks.container_instance import (
+ AzureContainerInstanceAsyncHook,
+ AzureContainerInstanceHook,
+)
@pytest.fixture
@@ -252,3 +260,256 @@ class TestAzureContainerInstanceHookCloudEnvironment:
base_url=expected_base_url,
credential_scopes=expected_scopes,
)
+
+
[email protected]
+def async_conn_with_credentials(create_mock_connection):
+ return create_mock_connection(
+ Connection(
+ conn_id="azure_aci_async_test",
+ conn_type="azure_container_instance",
+ login="client-id",
+ password="client-secret",
+ extra={
+ "tenantId": "tenant-id",
+ "subscriptionId": "subscription-id",
+ },
+ )
+ )
+
+
[email protected]
+def async_conn_without_credentials(create_mock_connection):
+ return create_mock_connection(
+ Connection(
+ conn_id="azure_aci_async_no_creds",
+ conn_type="azure_container_instance",
+ extra={"subscriptionId": "subscription-id"},
+ )
+ )
+
+
+class TestAzureContainerInstanceAsyncHook:
+ @patch(
+
"airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient"
+ )
+
@patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential")
+ @pytest.mark.asyncio
+ async def test_get_async_conn_with_client_secret(
+ self,
+ mock_credential_cls,
+ mock_client_cls,
+ async_conn_with_credentials,
+ ):
+ mock_credential = MagicMock(spec=AsyncClientSecretCredential)
+ mock_credential.close = AsyncMock()
+ mock_credential_cls.return_value = mock_credential
+ mock_client_instance =
MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_instance.close = AsyncMock()
+ mock_client_cls.return_value = mock_client_instance
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ async with hook.get_async_conn() as conn:
+ mock_credential_cls.assert_called_once_with(
+ client_id="client-id",
+ client_secret="client-secret",
+ tenant_id="tenant-id",
+ )
+ mock_client_cls.assert_called_once_with(
+ credential=mock_credential,
+ subscription_id="subscription-id",
+ )
+ assert conn is mock_client_instance
+
+ @patch(
+
"airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient"
+ )
+
@patch("airflow.providers.microsoft.azure.hooks.container_instance.get_async_default_azure_credential")
+ @pytest.mark.asyncio
+ async def test_get_async_conn_with_default_credential(
+ self,
+ mock_default_cred,
+ mock_client_cls,
+ async_conn_without_credentials,
+ ):
+ mock_credential = MagicMock()
+ mock_credential.close = AsyncMock()
+ mock_default_cred.return_value = mock_credential
+ mock_client_instance =
MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_instance.close = AsyncMock()
+ mock_client_cls.return_value = mock_client_instance
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_without_credentials.conn_id)
+ async with hook.get_async_conn() as conn:
+ mock_default_cred.assert_called_once_with(
+ managed_identity_client_id=None,
+ workload_identity_tenant_id=None,
+ )
+ assert conn is mock_client_instance
+
+ @patch(
+
"airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient"
+ )
+
@patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential")
+ @pytest.mark.asyncio
+ async def test_get_async_conn_closes_client_and_credential(
+ self,
+ mock_credential_cls,
+ mock_client_cls,
+ async_conn_with_credentials,
+ ):
+ mock_credential = AsyncMock(spec=AsyncClientSecretCredential)
+ mock_credential_cls.return_value = mock_credential
+ mock_client_instance =
AsyncMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ async with hook.get_async_conn() as conn:
+ assert conn is mock_client_instance
+
+ mock_client_instance.close.assert_called_once()
+ mock_credential.close.assert_called_once()
+
+ @patch(
+
"airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient"
+ )
+
@patch("airflow.providers.microsoft.azure.hooks.container_instance.get_async_default_azure_credential")
+ @pytest.mark.asyncio
+ async def test_get_async_conn_skips_credential_close_when_not_closeable(
+ self,
+ mock_default_cred,
+ mock_client_cls,
+ async_conn_without_credentials,
+ ):
+ # Credential without a close method (spec=[] removes all attributes)
+ mock_credential = MagicMock(spec=[])
+ mock_default_cred.return_value = mock_credential
+ mock_client_instance =
AsyncMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_without_credentials.conn_id)
+ async with hook.get_async_conn():
+ pass
+
+ mock_client_instance.close.assert_called_once()
+ assert not hasattr(mock_credential, "close")
+
+ @patch(
+
"airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient"
+ )
+
@patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential")
+ @pytest.mark.asyncio
+ async def test_get_async_conn_closes_on_exception(
+ self,
+ mock_credential_cls,
+ mock_client_cls,
+ async_conn_with_credentials,
+ ):
+ mock_credential = AsyncMock(spec=AsyncClientSecretCredential)
+ mock_credential_cls.return_value = mock_credential
+ mock_client_instance =
AsyncMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ with pytest.raises(RuntimeError, match="boom"):
+ async with hook.get_async_conn():
+ raise RuntimeError("boom")
+
+ mock_client_instance.close.assert_called_once()
+ mock_credential.close.assert_called_once()
+
+ @patch(
+
"airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient"
+ )
+
@patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential")
+ @pytest.mark.asyncio
+ async def test_get_state(
+ self,
+ mock_credential_cls,
+ mock_client_cls,
+ async_conn_with_credentials,
+ ):
+ mock_credential_cls.return_value =
AsyncMock(spec=AsyncClientSecretCredential)
+ mock_cg = MagicMock(spec=ContainerGroup)
+ mock_client = MagicMock()
+ mock_client.close = AsyncMock()
+ mock_client.container_groups.get = AsyncMock(return_value=mock_cg)
+ mock_client_cls.return_value = mock_client
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ result = await hook.get_state("my-rg", "my-container")
+
+ mock_client.container_groups.get.assert_called_once_with("my-rg",
"my-container")
+ assert result is mock_cg
+
+ @patch(
+
"airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient"
+ )
+
@patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential")
+ @pytest.mark.asyncio
+ async def test_get_logs(
+ self,
+ mock_credential_cls,
+ mock_client_cls,
+ async_conn_with_credentials,
+ ):
+ mock_credential_cls.return_value =
AsyncMock(spec=AsyncClientSecretCredential)
+ mock_logs = MagicMock(spec=Logs)
+ mock_logs.content = "line1\nline2\n"
+ mock_client = MagicMock()
+ mock_client.close = AsyncMock()
+ mock_client.containers.list_logs = AsyncMock(return_value=mock_logs)
+ mock_client_cls.return_value = mock_client
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ result = await hook.get_logs("my-rg", "my-container")
+
+ mock_client.containers.list_logs.assert_called_once_with(
+ "my-rg", "my-container", "my-container", tail=1000
+ )
+ assert result == ["line1\n", "line2\n"]
+
+ @patch(
+
"airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient"
+ )
+
@patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential")
+ @pytest.mark.asyncio
+ async def test_get_logs_returns_none_sentinel_when_content_is_none(
+ self,
+ mock_credential_cls,
+ mock_client_cls,
+ async_conn_with_credentials,
+ ):
+ mock_credential_cls.return_value =
AsyncMock(spec=AsyncClientSecretCredential)
+ mock_logs = MagicMock(spec=Logs)
+ mock_logs.content = None
+ mock_client = MagicMock()
+ mock_client.close = AsyncMock()
+ mock_client.containers.list_logs = AsyncMock(return_value=mock_logs)
+ mock_client_cls.return_value = mock_client
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ result = await hook.get_logs("my-rg", "my-container")
+ assert result == [None]
+
+ @patch(
+
"airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient"
+ )
+
@patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential")
+ @pytest.mark.asyncio
+ async def test_delete(
+ self,
+ mock_credential_cls,
+ mock_client_cls,
+ async_conn_with_credentials,
+ ):
+ mock_credential_cls.return_value =
AsyncMock(spec=AsyncClientSecretCredential)
+ mock_client = MagicMock()
+ mock_client.close = AsyncMock()
+ mock_client.container_groups.begin_delete = AsyncMock()
+ mock_client_cls.return_value = mock_client
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ await hook.delete("my-rg", "my-container")
+
+
mock_client.container_groups.begin_delete.assert_called_once_with("my-rg",
"my-container")
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py
index 9325b318bce..d21070de0c3 100644
---
a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py
@@ -32,8 +32,10 @@ from azure.mgmt.containerinstance.models import (
Event,
)
+from airflow.exceptions import TaskDeferred
from airflow.providers.common.compat.sdk import AirflowException
from airflow.providers.microsoft.azure.operators.container_instances import
AzureContainerInstancesOperator
+from airflow.providers.microsoft.azure.triggers.container_instance import
AzureContainerInstanceTrigger
from tests_common.test_utils.compat import Context
@@ -642,6 +644,292 @@ class TestACIOperator:
# 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_unconditionally(self, aci_mock):
+ """When deferrable=True, defer() is called immediately after
create_or_update."""
+ 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 aci_mock.return_value.create_or_update.call_count == 1
+ 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"
+ assert aci_mock.return_value.delete.call_count == 0
+ assert aci_mock.return_value.get_state.call_count == 0
+
+ @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 RuntimeError when event has status=error."""
+ 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_error=False,
+ )
+ with pytest.raises(RuntimeError, match="Container group failed"):
+ aci.execute_complete(
+ context=None,
+ event={
+ "status": "error",
+ "exit_code": 1,
+ "detail_status": "OOMKilled",
+ "message": "Container group failed with exit code 1",
+ "resource_group": "resource-group",
+ "name": "container-name",
+ },
+ )
+
+ @mock.patch(
+
"airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_complete_error_event_removes_on_error(self, aci_mock):
+ """execute_complete with remove_on_error=True deletes the container on
error."""
+ 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_error=True,
+ )
+ with pytest.raises(RuntimeError):
+ aci.execute_complete(
+ context=None,
+ event={
+ "status": "error",
+ "exit_code": 1,
+ "detail_status": "Failed",
+ "message": "Container group failed with exit code 1",
+ "resource_group": "resource-group",
+ "name": "container-name",
+ },
+ )
+ 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_none_event_raises(self, aci_mock):
+ """execute_complete raises ValueError when event is 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,
+ )
+ with pytest.raises(ValueError, match="event is None"):
+ aci.execute_complete(context=None, event=None)
+
+ @mock.patch(
+
"airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_complete_pushes_all_logs_to_xcom(self, aci_mock):
+ """execute_complete pushes all logs to XCom when xcom_all=True."""
+ aci_mock.return_value.get_logs.return_value = ["line1\n", "line2\n"]
+
+ 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,
+ xcom_all=True,
+ remove_on_success=False,
+ )
+ context = Context(ti=XcomMock())
+ aci.execute_complete(
+ context=context,
+ event={
+ "status": "success",
+ "exit_code": 0,
+ "detail_status": "Completed",
+ "resource_group": "resource-group",
+ "name": "container-name",
+ },
+ )
+ assert context["ti"].xcom_pull(key="logs") == ["line1\n", "line2\n"]
+
+ @mock.patch(
+
"airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_complete_pushes_last_log_to_xcom(self, aci_mock):
+ """execute_complete pushes only last log line to XCom when
xcom_all=False."""
+ aci_mock.return_value.get_logs.return_value = ["line1\n", "line2\n"]
+
+ 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,
+ xcom_all=False,
+ remove_on_success=False,
+ )
+ context = Context(ti=XcomMock())
+ aci.execute_complete(
+ context=context,
+ event={
+ "status": "success",
+ "exit_code": 0,
+ "detail_status": "Completed",
+ "resource_group": "resource-group",
+ "name": "container-name",
+ },
+ )
+ assert context["ti"].xcom_pull(key="logs") == ["line2\n"]
+
+ @mock.patch(
+
"airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_remove_on_success_false_sync_path(self, aci_mock):
+ """Non-deferrable: remove_on_success=False should NOT delete the
container on success."""
+ 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=None,
+ registry_conn_id=None,
+ resource_group="resource-group",
+ name="container-name",
+ image="container-image",
+ region="region",
+ task_id="task",
+ remove_on_success=False,
+ )
+ aci.execute(None)
+ assert aci_mock.return_value.delete.call_count == 0
+
+ @mock.patch(
+
"airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_remove_on_success_true_sync_path(self, aci_mock):
+ """Non-deferrable: remove_on_success=True (default) deletes the
container on success."""
+ 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=None,
+ registry_conn_id=None,
+ resource_group="resource-group",
+ name="container-name",
+ image="container-image",
+ region="region",
+ task_id="task",
+ remove_on_success=True,
+ )
+ aci.execute(None)
+ assert aci_mock.return_value.delete.call_count == 1
+
class XcomMock:
def __init__(self) -> None:
diff --git
a/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_container_instance.py
b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_container_instance.py
new file mode 100644
index 00000000000..e4d26060ed2
--- /dev/null
+++
b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_container_instance.py
@@ -0,0 +1,248 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from unittest import mock
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from azure.mgmt.containerinstance.models import (
+ Container,
+ ContainerGroup,
+ ContainerPropertiesInstanceView,
+ ContainerState,
+)
+
+from airflow.providers.microsoft.azure.triggers.container_instance import
AzureContainerInstanceTrigger
+from airflow.triggers.base import TriggerEvent
+
+RESOURCE_GROUP = "test-resource-group"
+CONTAINER_NAME = "test-container"
+CI_CONN_ID = "azure_container_instance_test"
+POLLING_INTERVAL = 5.0
+
+MODULE = "airflow.providers.microsoft.azure.triggers.container_instance"
+
+
+def _make_cg_state(state: str, exit_code: int = 0, detail_status: str =
"test") -> ContainerGroup:
+ """Build a minimal fake ContainerGroup using real SDK models (read-only
attrs use type: ignore)."""
+ container_state = ContainerState()
+ container_state.state = state # type: ignore[assignment]
+ container_state.exit_code = exit_code # type: ignore[assignment]
+ container_state.detail_status = detail_status # type: ignore[assignment]
+ instance_view = ContainerPropertiesInstanceView()
+ instance_view.current_state = container_state # type: ignore[assignment]
+ container = Container(name="test", image="test", resources="test") #
type: ignore[arg-type]
+ container.instance_view = instance_view # type: ignore[assignment]
+ cg = ContainerGroup(containers=[container], os_type="Linux")
+ cg.provisioning_state = state # type: ignore[assignment]
+ return cg
+
+
+def _make_provisioning_state(state: str) -> ContainerGroup:
+ """Build a ContainerGroup where instance_view is None (provisioning
phase)."""
+ container = Container(name="test", image="test", resources="test") #
type: ignore[arg-type]
+ container.instance_view = None # type: ignore[assignment]
+ cg = ContainerGroup(containers=[container], os_type="Linux")
+ cg.provisioning_state = state # type: ignore[assignment]
+ return cg
+
+
+def _make_mock_hook(mock_client):
+ """Create a mock hook whose get_async_conn yields the given mock client."""
+ mock_hook = MagicMock()
+
+ @asynccontextmanager
+ async def _get_async_conn():
+ yield mock_client
+
+ mock_hook.get_async_conn = _get_async_conn
+ return mock_hook
+
+
+class TestAzureContainerInstanceTrigger:
+ def test_serialize(self):
+ trigger = AzureContainerInstanceTrigger(
+ resource_group=RESOURCE_GROUP,
+ name=CONTAINER_NAME,
+ ci_conn_id=CI_CONN_ID,
+ polling_interval=POLLING_INTERVAL,
+ )
+ classpath, kwargs = trigger.serialize()
+
+ assert classpath == (
+
"airflow.providers.microsoft.azure.triggers.container_instance.AzureContainerInstanceTrigger"
+ )
+ assert kwargs == {
+ "resource_group": RESOURCE_GROUP,
+ "name": CONTAINER_NAME,
+ "ci_conn_id": CI_CONN_ID,
+ "polling_interval": POLLING_INTERVAL,
+ }
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook")
+ async def test_run_yields_success_on_terminated_with_exit_code_zero(self,
mock_hook_cls):
+ mock_client = MagicMock()
+ mock_client.container_groups.get =
AsyncMock(return_value=_make_cg_state("Terminated", exit_code=0))
+ mock_hook_cls.return_value = _make_mock_hook(mock_client)
+
+ trigger = AzureContainerInstanceTrigger(
+ resource_group=RESOURCE_GROUP,
+ name=CONTAINER_NAME,
+ ci_conn_id=CI_CONN_ID,
+ polling_interval=POLLING_INTERVAL,
+ )
+ events = [event async for event in trigger.run()]
+
+ assert len(events) == 1
+ assert isinstance(events[0], TriggerEvent)
+ assert events[0].payload["status"] == "success"
+ assert events[0].payload["exit_code"] == 0
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook")
+ async def test_run_yields_success_on_succeeded_state(self, mock_hook_cls):
+ mock_client = MagicMock()
+ mock_client.container_groups.get =
AsyncMock(return_value=_make_cg_state("Succeeded", exit_code=0))
+ mock_hook_cls.return_value = _make_mock_hook(mock_client)
+
+ trigger = AzureContainerInstanceTrigger(
+ resource_group=RESOURCE_GROUP,
+ name=CONTAINER_NAME,
+ ci_conn_id=CI_CONN_ID,
+ )
+ events = [event async for event in trigger.run()]
+
+ assert events[0].payload["status"] == "success"
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook")
+ async def test_run_yields_error_on_terminated_with_nonzero_exit_code(self,
mock_hook_cls):
+ mock_client = MagicMock()
+ mock_client.container_groups.get =
AsyncMock(return_value=_make_cg_state("Terminated", exit_code=1))
+ mock_hook_cls.return_value = _make_mock_hook(mock_client)
+
+ trigger = AzureContainerInstanceTrigger(
+ resource_group=RESOURCE_GROUP,
+ name=CONTAINER_NAME,
+ ci_conn_id=CI_CONN_ID,
+ )
+ events = [event async for event in trigger.run()]
+
+ assert events[0].payload["status"] == "error"
+ assert events[0].payload["exit_code"] == 1
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook")
+ @pytest.mark.parametrize("terminal_state", ["Failed", "Unhealthy"])
+ async def test_run_yields_error_on_failed_states(self, mock_hook_cls,
terminal_state):
+ mock_client = MagicMock()
+ mock_client.container_groups.get =
AsyncMock(return_value=_make_cg_state(terminal_state, exit_code=1))
+ mock_hook_cls.return_value = _make_mock_hook(mock_client)
+
+ trigger = AzureContainerInstanceTrigger(
+ resource_group=RESOURCE_GROUP,
+ name=CONTAINER_NAME,
+ ci_conn_id=CI_CONN_ID,
+ )
+ events = [event async for event in trigger.run()]
+
+ assert events[0].payload["status"] == "error"
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+ @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook")
+ async def test_run_polls_through_non_terminal_states(self, mock_hook_cls,
mock_sleep):
+ mock_client = MagicMock()
+ mock_client.container_groups.get = AsyncMock(
+ side_effect=[
+ _make_provisioning_state("Creating"),
+ _make_cg_state("Running", exit_code=0),
+ _make_cg_state("Terminated", exit_code=0),
+ ]
+ )
+ mock_hook_cls.return_value = _make_mock_hook(mock_client)
+
+ trigger = AzureContainerInstanceTrigger(
+ resource_group=RESOURCE_GROUP,
+ name=CONTAINER_NAME,
+ ci_conn_id=CI_CONN_ID,
+ polling_interval=POLLING_INTERVAL,
+ )
+ events = [event async for event in trigger.run()]
+
+ assert mock_client.container_groups.get.call_count == 3
+ assert mock_sleep.call_count == 2
+ mock_sleep.assert_called_with(POLLING_INTERVAL)
+ assert events[0].payload["status"] == "success"
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook")
+ async def test_run_yields_error_event_on_exception(self, mock_hook_cls):
+ mock_client = MagicMock()
+ mock_client.container_groups.get =
AsyncMock(side_effect=Exception("Azure API error"))
+ mock_hook_cls.return_value = _make_mock_hook(mock_client)
+
+ trigger = AzureContainerInstanceTrigger(
+ resource_group=RESOURCE_GROUP,
+ name=CONTAINER_NAME,
+ ci_conn_id=CI_CONN_ID,
+ )
+ events = [event async for event in trigger.run()]
+
+ assert len(events) == 1
+ assert events[0].payload["status"] == "error"
+ assert "Azure API error" in events[0].payload["message"]
+
+ @pytest.mark.asyncio
+ @mock.patch(f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+ @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook")
+ async def test_run_provisioning_state_does_not_yield_until_terminal(self,
mock_hook_cls, mock_sleep):
+ """Container in 'Creating' (no instance_view) should keep polling."""
+ mock_client = MagicMock()
+ mock_client.container_groups.get = AsyncMock(
+ side_effect=[
+ _make_provisioning_state("Creating"),
+ _make_cg_state("Terminated", exit_code=0),
+ ]
+ )
+ mock_hook_cls.return_value = _make_mock_hook(mock_client)
+
+ trigger = AzureContainerInstanceTrigger(
+ resource_group=RESOURCE_GROUP,
+ name=CONTAINER_NAME,
+ ci_conn_id=CI_CONN_ID,
+ )
+ events = [event async for event in trigger.run()]
+
+ assert mock_client.container_groups.get.call_count == 2
+ assert events[0].payload["status"] == "success"
+
+ def test_serialize_deserialize_roundtrip(self):
+ original = AzureContainerInstanceTrigger(
+ resource_group=RESOURCE_GROUP,
+ name=CONTAINER_NAME,
+ ci_conn_id=CI_CONN_ID,
+ polling_interval=10.0,
+ )
+ classpath, kwargs = original.serialize()
+ reconstructed = AzureContainerInstanceTrigger(**kwargs)
+ _, kwargs2 = reconstructed.serialize()
+ assert kwargs == kwargs2
