From 5f2c6f37725523f8c0edae85475b811e3c064afd Mon Sep 17 00:00:00 2001 From: Pankaj Date: Thu, 24 Aug 2023 20:00:49 +0530 Subject: [PATCH 1/3] Bump azure-mgmt-containerinstance>=7.0.0,<9.0.0 --- .../azure/hooks/container_instance.py | 68 +++++++++- .../azure/operators/container_instances.py | 3 +- .../providers/microsoft/azure/provider.yaml | 2 +- generated/provider_dependencies.json | 2 +- .../hooks/test_azure_container_instance.py | 6 +- .../test_azure_container_instances.py | 120 +++++++++++------- 6 files changed, 144 insertions(+), 57 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/container_instance.py b/airflow/providers/microsoft/azure/hooks/container_instance.py index 9a0d0ec21067f..c2a399a47f457 100644 --- a/airflow/providers/microsoft/azure/hooks/container_instance.py +++ b/airflow/providers/microsoft/azure/hooks/container_instance.py @@ -18,11 +18,15 @@ from __future__ import annotations import warnings +from typing import Any +from azure.common.client_factory import get_client_from_auth_file, get_client_from_json_dict +from azure.common.credentials import ServicePrincipalCredentials +from azure.identity import DefaultAzureCredential from azure.mgmt.containerinstance import ContainerInstanceManagementClient from azure.mgmt.containerinstance.models import ContainerGroup -from airflow.exceptions import AirflowProviderDeprecationWarning +from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.microsoft.azure.hooks.base_azure import AzureBaseHook @@ -49,6 +53,59 @@ def __init__(self, azure_conn_id: str = default_conn_name) -> None: super().__init__(sdk_client=ContainerInstanceManagementClient, conn_id=azure_conn_id) self.connection = self.get_conn() + def get_conn(self) -> Any: + """ + Authenticates the resource using the connection id passed during init. + + :return: the authenticated client. + """ + conn = self.get_connection(self.conn_id) + tenant = conn.extra_dejson.get("tenantId") + if not tenant and conn.extra_dejson.get("extra__azure__tenantId"): + warnings.warn( + "`extra__azure__tenantId` is deprecated in azure connection extra, " + "please use `tenantId` instead", + AirflowProviderDeprecationWarning, + stacklevel=2, + ) + tenant = conn.extra_dejson.get("extra__azure__tenantId") + subscription_id = conn.extra_dejson.get("subscriptionId") + if not subscription_id and conn.extra_dejson.get("extra__azure__subscriptionId"): + warnings.warn( + "`extra__azure__subscriptionId` is deprecated in azure connection extra, " + "please use `subscriptionId` instead", + AirflowProviderDeprecationWarning, + stacklevel=2, + ) + subscription_id = conn.extra_dejson.get("extra__azure__subscriptionId") + + key_path = conn.extra_dejson.get("key_path") + if key_path: + if not key_path.endswith(".json"): + raise AirflowException("Unrecognised extension for key file.") + self.log.info("Getting connection using a JSON key file.") + return get_client_from_auth_file(client_class=self.sdk_client, auth_path=key_path) + + key_json = conn.extra_dejson.get("key_json") + if key_json: + self.log.info("Getting connection using a JSON config.") + return get_client_from_json_dict(client_class=self.sdk_client, config_dict=key_json) + + credential: ServicePrincipalCredentials | DefaultAzureCredential + if all([conn.login, conn.password, tenant]): + self.log.info("Getting connection using specific credentials and subscription_id.") + credential = ServicePrincipalCredentials( + client_id=conn.login, secret=conn.password, tenant=tenant + ) + else: + self.log.info("Using DefaultAzureCredential as credential") + credential = DefaultAzureCredential() + + return ContainerInstanceManagementClient( + credential=credential, + subscription_id=subscription_id, + ) + def create_or_update(self, resource_group: str, name: str, container_group: ContainerGroup) -> None: """ Create a new container group. @@ -57,7 +114,7 @@ def create_or_update(self, resource_group: str, name: str, container_group: Cont :param name: the name of the container group :param container_group: the properties of the container group """ - self.connection.container_groups.create_or_update(resource_group, name, container_group) + self.connection.container_groups.begin_create_or_update(resource_group, name, container_group) def get_state_exitcode_details(self, resource_group: str, name: str) -> tuple: """ @@ -113,7 +170,7 @@ def get_logs(self, resource_group: str, name: str, tail: int = 1000) -> list: :param tail: the size of the tail :return: A list of log messages """ - logs = self.connection.container.list_logs(resource_group, name, name, tail=tail) + logs = self.connection.containers.list_logs(resource_group, name, name, tail=tail) return logs.content.splitlines(True) def delete(self, resource_group: str, name: str) -> None: @@ -123,7 +180,7 @@ def delete(self, resource_group: str, name: str) -> None: :param resource_group: the name of the resource group :param name: the name of the container group """ - self.connection.container_groups.delete(resource_group, name) + self.connection.container_groups.begin_delete(resource_group, name) def exists(self, resource_group: str, name: str) -> bool: """ @@ -137,6 +194,9 @@ def exists(self, resource_group: str, name: str) -> bool: return True return False + def list_container_group(self, resource_group: str, name: str): + return self.connection.container_groups.list(resource_group, name) + def test_connection(self): """Test a configured Azure Container Instance connection.""" try: diff --git a/airflow/providers/microsoft/azure/operators/container_instances.py b/airflow/providers/microsoft/azure/operators/container_instances.py index 8b5fba26cdcb1..79b20c5c80c9e 100644 --- a/airflow/providers/microsoft/azure/operators/container_instances.py +++ b/airflow/providers/microsoft/azure/operators/container_instances.py @@ -293,9 +293,8 @@ def _monitor_logging(self, resource_group: str, name: str) -> int: while True: try: - cg_state = self._ci_hook.get_state(resource_group, name) + cg_state = self._ci_hook.list_container_group(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: c_state = instance_view.current_state diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index ab713bdbe1d84..a18e52cfd7622 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -79,9 +79,9 @@ dependencies: - adal>=1.2.7 - azure-storage-file-datalake>=12.9.1 - azure-kusto-data>=4.1.0 + - azure-mgmt-containerinstance>=7.0.0,<9.0.0 # TODO: upgrade to newer versions of all the below libraries. # See issue https://github.com/apache/airflow/issues/30199 - - azure-mgmt-containerinstance>=1.5.0,<2.0 - azure-mgmt-datafactory>=1.0.0,<2.0 integrations: diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 5191600a98392..98c5e4e281834 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -559,7 +559,7 @@ "azure-identity>=1.3.1", "azure-keyvault-secrets>=4.1.0", "azure-kusto-data>=4.1.0", - "azure-mgmt-containerinstance>=1.5.0,<2.0", + "azure-mgmt-containerinstance>=7.0.0,<9.0.0", "azure-mgmt-datafactory>=1.0.0,<2.0", "azure-mgmt-datalake-store>=0.5.0", "azure-mgmt-resource>=2.2.0", diff --git a/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py b/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py index a10dde9c651ae..ab454bf9868f3 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py @@ -53,7 +53,7 @@ def setup_method(self): self.hook = AzureContainerInstanceHook(azure_conn_id="azure_container_instance_test") @patch("azure.mgmt.containerinstance.models.ContainerGroup") - @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.create_or_update") + @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.begin_create_or_update") def test_create_or_update(self, create_or_update_mock, container_group_mock): self.hook.create_or_update("resource_group", "aci-test", container_group_mock) create_or_update_mock.assert_called_once_with("resource_group", "aci-test", container_group_mock) @@ -63,7 +63,7 @@ def test_get_state(self, get_state_mock): self.hook.get_state("resource_group", "aci-test") get_state_mock.assert_called_once_with("resource_group", "aci-test", raw=False) - @patch("azure.mgmt.containerinstance.operations.ContainerOperations.list_logs") + @patch("azure.mgmt.containerinstance.operations.ContainersOperations.list_logs") def test_get_logs(self, list_logs_mock): expected_messages = ["log line 1\n", "log line 2\n", "log line 3\n"] logs = Logs(content="".join(expected_messages)) @@ -73,7 +73,7 @@ def test_get_logs(self, list_logs_mock): assert logs == expected_messages - @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.delete") + @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.begin_delete") def test_delete(self, delete_mock): self.hook.delete("resource_group", "aci-test") delete_mock.assert_called_once_with("resource_group", "aci-test") diff --git a/tests/providers/microsoft/azure/operators/test_azure_container_instances.py b/tests/providers/microsoft/azure/operators/test_azure_container_instances.py index 6c7cd1dbb5009..a83e65e4189bd 100644 --- a/tests/providers/microsoft/azure/operators/test_azure_container_instances.py +++ b/tests/providers/microsoft/azure/operators/test_azure_container_instances.py @@ -22,7 +22,13 @@ from unittest.mock import MagicMock import pytest -from azure.mgmt.containerinstance.models import ContainerState, Event +from azure.mgmt.containerinstance.models import ( + Container, + ContainerGroup, + ContainerPropertiesInstanceView, + ContainerState, + Event, +) from airflow.exceptions import AirflowException from airflow.providers.microsoft.azure.operators.container_instances import AzureContainerInstancesOperator @@ -35,10 +41,12 @@ def make_mock_cg(container_state, events=None): """ events = events or [] instance_view_dict = {"current_state": container_state, "events": events} - instance_view = namedtuple("InstanceView", instance_view_dict.keys())(*instance_view_dict.values()) + instance_view = namedtuple("ContainerPropertiesInstanceView", instance_view_dict.keys())( + *instance_view_dict.values() + ) container_dict = {"instance_view": instance_view} - container = namedtuple("Container", container_dict.keys())(*container_dict.values()) + container = namedtuple("Containers", container_dict.keys())(*container_dict.values()) container_g_dict = {"containers": [container]} container_g = namedtuple("ContainerGroup", container_g_dict.keys())(*container_g_dict.values()) @@ -53,23 +61,42 @@ def make_mock_cg_with_missing_events(container_state): This can happen, when the container group is provisioned, but not started. """ instance_view_dict = {"current_state": container_state, "events": None} - instance_view = namedtuple("InstanceView", instance_view_dict.keys())(*instance_view_dict.values()) + instance_view = namedtuple("ContainerPropertiesInstanceView", instance_view_dict.keys())( + *instance_view_dict.values() + ) container_dict = {"instance_view": instance_view} - container = namedtuple("Container", container_dict.keys())(*container_dict.values()) + container = namedtuple("Containers", container_dict.keys())(*container_dict.values()) container_g_dict = {"containers": [container]} container_g = namedtuple("ContainerGroup", container_g_dict.keys())(*container_g_dict.values()) return container_g +def make_mock_container(state: str, exit_code: int, detail_status: str, events: Event | None = None): + container = Container(name="hello_world", image="test", resources="test") + container_prop = ContainerPropertiesInstanceView() + container_state = ContainerState() + container_state.state = state + container_state.exit_code = exit_code + container_state.detail_status = detail_status + container_prop.current_state = container_state + if events: + container_prop.events = events + container.instance_view = container_prop + + cg = ContainerGroup(containers=[container], os_type="Linux") + + return cg + + class TestACIOperator: @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute(self, aci_mock): - expected_c_state = ContainerState(state="Terminated", exit_code=0, detail_status="test") - expected_cg = make_mock_cg(expected_c_state) + expected_cg = make_mock_container(state="Terminated", exit_code=0, detail_status="test") + + aci_mock.return_value.list_container_group.return_value = expected_cg - aci_mock.return_value.get_state.return_value = expected_cg aci_mock.return_value.exists.return_value = False aci = AzureContainerInstancesOperator( @@ -102,10 +129,10 @@ def test_execute(self, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute_with_failures(self, aci_mock): - expected_c_state = ContainerState(state="Terminated", exit_code=1, detail_status="test") - expected_cg = make_mock_cg(expected_c_state) - aci_mock.return_value.get_state.return_value = expected_cg + expected_cg = make_mock_container(state="Terminated", exit_code=1, detail_status="test") + aci_mock.return_value.list_container_group.return_value = expected_cg + aci_mock.return_value.exists.return_value = False aci = AzureContainerInstancesOperator( @@ -124,11 +151,11 @@ def test_execute_with_failures(self, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute_with_tags(self, aci_mock): - expected_c_state = ContainerState(state="Terminated", exit_code=0, detail_status="test") - expected_cg = make_mock_cg(expected_c_state) + + expected_cg = make_mock_container(state="Terminated", exit_code=0, detail_status="test") + aci_mock.return_value.list_container_group.return_value = expected_cg tags = {"testKey": "testValue"} - aci_mock.return_value.get_state.return_value = expected_cg aci_mock.return_value.exists.return_value = False aci = AzureContainerInstancesOperator( @@ -163,15 +190,20 @@ def test_execute_with_tags(self, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute_with_messages_logs(self, aci_mock): - events = [Event(message="test"), Event(message="messages")] - expected_c_state1 = ContainerState(state="Succeeded", exit_code=0, detail_status="test") - expected_cg1 = make_mock_cg(expected_c_state1, events) - expected_c_state2 = ContainerState(state="Running", exit_code=0, detail_status="test") - expected_cg2 = make_mock_cg(expected_c_state2, events) - expected_c_state3 = ContainerState(state="Terminated", exit_code=0, detail_status="test") - expected_cg3 = make_mock_cg(expected_c_state3, events) - - aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg2, expected_cg3] + event1 = Event() + event1.message = "test" + event2 = Event() + event2.message = "messages" + events = [event1, event2] + expected_cg1 = make_mock_container( + state="Succeeded", exit_code=0, detail_status="test", events=events + ) + expected_cg2 = make_mock_container(state="Running", exit_code=0, detail_status="test", events=events) + expected_cg3 = make_mock_container( + state="Terminated", exit_code=0, detail_status="test", events=events + ) + + aci_mock.return_value.list_container_group.side_effect = [expected_cg1, expected_cg2, expected_cg3] aci_mock.return_value.get_logs.return_value = ["test", "logs"] aci_mock.return_value.exists.return_value = False @@ -187,7 +219,7 @@ def test_execute_with_messages_logs(self, aci_mock): aci.execute(None) assert aci_mock.return_value.create_or_update.call_count == 1 - assert aci_mock.return_value.get_state.call_count == 3 + assert aci_mock.return_value.list_container_group.call_count == 3 assert aci_mock.return_value.get_logs.call_count == 3 assert aci_mock.return_value.delete.call_count == 1 @@ -211,11 +243,11 @@ def test_name_checker(self): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute_with_ipaddress(self, aci_mock): - expected_c_state = ContainerState(state="Terminated", exit_code=0, detail_status="test") - expected_cg = make_mock_cg(expected_c_state) ipaddress = MagicMock() - aci_mock.return_value.get_state.return_value = expected_cg + aci_mock.return_value.list_container_group.return_value = make_mock_container( + state="Terminated", exit_code=0, detail_status="test" + ) aci_mock.return_value.exists.return_value = False aci = AzureContainerInstancesOperator( @@ -236,10 +268,10 @@ def test_execute_with_ipaddress(self, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute_with_windows_os_and_diff_restart_policy(self, aci_mock): - expected_c_state = ContainerState(state="Terminated", exit_code=0, detail_status="test") - expected_cg = make_mock_cg(expected_c_state) - aci_mock.return_value.get_state.return_value = expected_cg + aci_mock.return_value.list_container_group.return_value = make_mock_container( + state="Terminated", exit_code=0, detail_status="test" + ) aci_mock.return_value.exists.return_value = False aci = AzureContainerInstancesOperator( @@ -262,10 +294,10 @@ def test_execute_with_windows_os_and_diff_restart_policy(self, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute_fails_with_incorrect_os_type(self, aci_mock): - expected_c_state = ContainerState(state="Terminated", exit_code=0, detail_status="test") - expected_cg = make_mock_cg(expected_c_state) - aci_mock.return_value.get_state.return_value = expected_cg + aci_mock.return_value.list_container_group.return_value = make_mock_container( + state="Terminated", exit_code=0, detail_status="test" + ) aci_mock.return_value.exists.return_value = False with pytest.raises(AirflowException) as ctx: @@ -288,10 +320,10 @@ def test_execute_fails_with_incorrect_os_type(self, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute_fails_with_incorrect_restart_policy(self, aci_mock): - expected_c_state = ContainerState(state="Terminated", exit_code=0, detail_status="test") - expected_cg = make_mock_cg(expected_c_state) - aci_mock.return_value.get_state.return_value = expected_cg + aci_mock.return_value.list_container_group.return_value = make_mock_container( + state="Terminated", exit_code=0, detail_status="test" + ) aci_mock.return_value.exists.return_value = False with pytest.raises(AirflowException) as ctx: @@ -315,12 +347,10 @@ def test_execute_fails_with_incorrect_restart_policy(self, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.sleep") def test_execute_correct_sleep_cycle(self, sleep_mock, aci_mock): - expected_c_state1 = ContainerState(state="Running", exit_code=0, detail_status="test") - expected_cg1 = make_mock_cg(expected_c_state1) - expected_c_state2 = ContainerState(state="Terminated", exit_code=0, detail_status="test") - expected_cg2 = make_mock_cg(expected_c_state2) + expected_cg1 = make_mock_container(state="Running", exit_code=0, detail_status="test") + expected_cg2 = make_mock_container(state="Terminated", exit_code=0, detail_status="test") - aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg1, expected_cg2] + aci_mock.return_value.list_container_group.side_effect = [expected_cg1, expected_cg1, expected_cg2] aci_mock.return_value.exists.return_value = False aci = AzureContainerInstancesOperator( @@ -340,12 +370,10 @@ def test_execute_correct_sleep_cycle(self, sleep_mock, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") @mock.patch("logging.Logger.exception") def test_execute_with_missing_events(self, log_mock, aci_mock): - expected_c_state1 = ContainerState(state="Running", exit_code=0, detail_status="test") - expected_cg1 = make_mock_cg_with_missing_events(expected_c_state1) - expected_c_state2 = ContainerState(state="Terminated", exit_code=0, detail_status="test") - expected_cg2 = make_mock_cg(expected_c_state2) + expected_cg1 = make_mock_container(state="Running", exit_code=0, detail_status="test") + expected_cg2 = make_mock_container(state="Terminated", exit_code=0, detail_status="test") - aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg2] + aci_mock.return_value.list_container_group.side_effect = [expected_cg1, expected_cg2] aci_mock.return_value.exists.return_value = False aci = AzureContainerInstancesOperator( From 04f92b03880b89a0658d670a863dbc1ea4adb63a Mon Sep 17 00:00:00 2001 From: Pankaj Date: Mon, 28 Aug 2023 13:52:13 +0530 Subject: [PATCH 2/3] Apply review suggestions --- .../azure/hooks/container_instance.py | 5 +---- .../azure/operators/container_instances.py | 4 ++-- .../providers/microsoft/azure/provider.yaml | 2 +- .../hooks/test_azure_container_instance.py | 2 +- .../test_azure_container_instances.py | 22 +++++++++---------- 5 files changed, 16 insertions(+), 19 deletions(-) diff --git a/airflow/providers/microsoft/azure/hooks/container_instance.py b/airflow/providers/microsoft/azure/hooks/container_instance.py index c2a399a47f457..8dbf1d0b58824 100644 --- a/airflow/providers/microsoft/azure/hooks/container_instance.py +++ b/airflow/providers/microsoft/azure/hooks/container_instance.py @@ -159,7 +159,7 @@ def get_state(self, resource_group: str, name: str) -> ContainerGroup: :param name: the name of the container group :return: ContainerGroup """ - return self.connection.container_groups.get(resource_group, name, raw=False) + return self.connection.container_groups.get(resource_group, name) def get_logs(self, resource_group: str, name: str, tail: int = 1000) -> list: """ @@ -194,9 +194,6 @@ def exists(self, resource_group: str, name: str) -> bool: return True return False - def list_container_group(self, resource_group: str, name: str): - return self.connection.container_groups.list(resource_group, name) - def test_connection(self): """Test a configured Azure Container Instance connection.""" try: diff --git a/airflow/providers/microsoft/azure/operators/container_instances.py b/airflow/providers/microsoft/azure/operators/container_instances.py index 79b20c5c80c9e..efbe8259b101a 100644 --- a/airflow/providers/microsoft/azure/operators/container_instances.py +++ b/airflow/providers/microsoft/azure/operators/container_instances.py @@ -126,11 +126,11 @@ def __init__( self, *, ci_conn_id: str, - registry_conn_id: str | None, resource_group: str, name: str, image: str, region: str, + registry_conn_id: str | None = None, environment_variables: dict | None = None, secured_variables: str | None = None, volumes: list | None = None, @@ -293,7 +293,7 @@ def _monitor_logging(self, resource_group: str, name: str) -> int: while True: try: - cg_state = self._ci_hook.list_container_group(resource_group, name) + cg_state = self._ci_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: diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index a18e52cfd7622..1be78ca6ccf34 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -79,9 +79,9 @@ dependencies: - adal>=1.2.7 - azure-storage-file-datalake>=12.9.1 - azure-kusto-data>=4.1.0 - - azure-mgmt-containerinstance>=7.0.0,<9.0.0 # TODO: upgrade to newer versions of all the below libraries. # See issue https://github.com/apache/airflow/issues/30199 + - azure-mgmt-containerinstance>=7.0.0,<9.0.0 - azure-mgmt-datafactory>=1.0.0,<2.0 integrations: diff --git a/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py b/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py index ab454bf9868f3..966cfcf30dc93 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_container_instance.py @@ -61,7 +61,7 @@ def test_create_or_update(self, create_or_update_mock, container_group_mock): @patch("azure.mgmt.containerinstance.operations.ContainerGroupsOperations.get") def test_get_state(self, get_state_mock): self.hook.get_state("resource_group", "aci-test") - get_state_mock.assert_called_once_with("resource_group", "aci-test", raw=False) + get_state_mock.assert_called_once_with("resource_group", "aci-test") @patch("azure.mgmt.containerinstance.operations.ContainersOperations.list_logs") def test_get_logs(self, list_logs_mock): diff --git a/tests/providers/microsoft/azure/operators/test_azure_container_instances.py b/tests/providers/microsoft/azure/operators/test_azure_container_instances.py index a83e65e4189bd..2984fe60eeada 100644 --- a/tests/providers/microsoft/azure/operators/test_azure_container_instances.py +++ b/tests/providers/microsoft/azure/operators/test_azure_container_instances.py @@ -95,7 +95,7 @@ class TestACIOperator: def test_execute(self, aci_mock): expected_cg = make_mock_container(state="Terminated", exit_code=0, detail_status="test") - aci_mock.return_value.list_container_group.return_value = expected_cg + aci_mock.return_value.get_state.return_value = expected_cg aci_mock.return_value.exists.return_value = False @@ -131,7 +131,7 @@ def test_execute(self, aci_mock): def test_execute_with_failures(self, aci_mock): expected_cg = make_mock_container(state="Terminated", exit_code=1, detail_status="test") - aci_mock.return_value.list_container_group.return_value = expected_cg + aci_mock.return_value.get_state.return_value = expected_cg aci_mock.return_value.exists.return_value = False @@ -153,7 +153,7 @@ def test_execute_with_failures(self, aci_mock): def test_execute_with_tags(self, aci_mock): expected_cg = make_mock_container(state="Terminated", exit_code=0, detail_status="test") - aci_mock.return_value.list_container_group.return_value = expected_cg + aci_mock.return_value.get_state.return_value = expected_cg tags = {"testKey": "testValue"} aci_mock.return_value.exists.return_value = False @@ -203,7 +203,7 @@ def test_execute_with_messages_logs(self, aci_mock): state="Terminated", exit_code=0, detail_status="test", events=events ) - aci_mock.return_value.list_container_group.side_effect = [expected_cg1, expected_cg2, expected_cg3] + aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg2, expected_cg3] aci_mock.return_value.get_logs.return_value = ["test", "logs"] aci_mock.return_value.exists.return_value = False @@ -219,7 +219,7 @@ def test_execute_with_messages_logs(self, aci_mock): aci.execute(None) assert aci_mock.return_value.create_or_update.call_count == 1 - assert aci_mock.return_value.list_container_group.call_count == 3 + assert aci_mock.return_value.get_state.call_count == 3 assert aci_mock.return_value.get_logs.call_count == 3 assert aci_mock.return_value.delete.call_count == 1 @@ -245,7 +245,7 @@ def test_name_checker(self): def test_execute_with_ipaddress(self, aci_mock): ipaddress = MagicMock() - aci_mock.return_value.list_container_group.return_value = make_mock_container( + aci_mock.return_value.get_state.return_value = make_mock_container( state="Terminated", exit_code=0, detail_status="test" ) aci_mock.return_value.exists.return_value = False @@ -269,7 +269,7 @@ def test_execute_with_ipaddress(self, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute_with_windows_os_and_diff_restart_policy(self, aci_mock): - aci_mock.return_value.list_container_group.return_value = make_mock_container( + aci_mock.return_value.get_state.return_value = make_mock_container( state="Terminated", exit_code=0, detail_status="test" ) aci_mock.return_value.exists.return_value = False @@ -295,7 +295,7 @@ def test_execute_with_windows_os_and_diff_restart_policy(self, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute_fails_with_incorrect_os_type(self, aci_mock): - aci_mock.return_value.list_container_group.return_value = make_mock_container( + aci_mock.return_value.get_state.return_value = make_mock_container( state="Terminated", exit_code=0, detail_status="test" ) aci_mock.return_value.exists.return_value = False @@ -321,7 +321,7 @@ def test_execute_fails_with_incorrect_os_type(self, aci_mock): @mock.patch("airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook") def test_execute_fails_with_incorrect_restart_policy(self, aci_mock): - aci_mock.return_value.list_container_group.return_value = make_mock_container( + aci_mock.return_value.get_state.return_value = make_mock_container( state="Terminated", exit_code=0, detail_status="test" ) aci_mock.return_value.exists.return_value = False @@ -350,7 +350,7 @@ def test_execute_correct_sleep_cycle(self, sleep_mock, aci_mock): expected_cg1 = make_mock_container(state="Running", exit_code=0, detail_status="test") expected_cg2 = make_mock_container(state="Terminated", exit_code=0, detail_status="test") - aci_mock.return_value.list_container_group.side_effect = [expected_cg1, expected_cg1, expected_cg2] + aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg1, expected_cg2] aci_mock.return_value.exists.return_value = False aci = AzureContainerInstancesOperator( @@ -373,7 +373,7 @@ def test_execute_with_missing_events(self, log_mock, aci_mock): expected_cg1 = make_mock_container(state="Running", exit_code=0, detail_status="test") expected_cg2 = make_mock_container(state="Terminated", exit_code=0, detail_status="test") - aci_mock.return_value.list_container_group.side_effect = [expected_cg1, expected_cg2] + aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg2] aci_mock.return_value.exists.return_value = False aci = AzureContainerInstancesOperator( From 16b92fc115c87cb6f2c64e06e1a6a9d7d8067884 Mon Sep 17 00:00:00 2001 From: Pankaj Singh <98807258+pankajastro@users.noreply.github.com> Date: Mon, 28 Aug 2023 15:00:10 +0530 Subject: [PATCH 3/3] Update airflow/providers/microsoft/azure/hooks/container_instance.py --- airflow/providers/microsoft/azure/hooks/container_instance.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/microsoft/azure/hooks/container_instance.py b/airflow/providers/microsoft/azure/hooks/container_instance.py index 1ca2091f9a098..cdebe7a241acc 100644 --- a/airflow/providers/microsoft/azure/hooks/container_instance.py +++ b/airflow/providers/microsoft/azure/hooks/container_instance.py @@ -19,7 +19,7 @@ import warnings from functools import cached_property -from typing import Any, TYPE_CHECKING +from typing import TYPE_CHECKING, Any from azure.common.client_factory import get_client_from_auth_file, get_client_from_json_dict from azure.common.credentials import ServicePrincipalCredentials