cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2899310869
##########
providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py:
##########
@@ -160,3 +167,177 @@ def test_get_conn_fallback_to_default_azure_credential(
credential=mock_credential,
subscription_id="subscription_id",
)
+
+
[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_cls.return_value = mock_credential
+ mock_client_instance =
MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ conn = await hook.get_async_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 == 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(spec=AsyncClientSecretCredential)
+ mock_default_cred.return_value = mock_credential
+ mock_client_instance =
MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_without_credentials.conn_id)
+ conn = await hook.get_async_conn()
+
+ mock_default_cred.assert_called_once_with(
+ managed_identity_client_id=None,
+ workload_identity_tenant_id=None,
+ )
+ assert conn == mock_client_instance
+
+ @pytest.mark.asyncio
+ async def test_get_async_conn_returns_cached_connection(self,
async_conn_with_credentials):
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_conn = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ hook._async_conn = mock_conn
+
+ conn = await hook.get_async_conn()
+ assert conn is mock_conn
+
+ @pytest.mark.asyncio
+ async def test_get_state(self, async_conn_with_credentials):
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_cg = MagicMock(spec=ContainerGroup)
+ mock_client.container_groups.get = AsyncMock(return_value=mock_cg)
+ hook._async_conn = mock_client
+
+ 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
+
+ @pytest.mark.asyncio
+ async def test_get_logs(self, async_conn_with_credentials):
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_logs = MagicMock(spec=Logs)
+ mock_logs.content = "line1\nline2\n"
+ mock_client.containers.list_logs = AsyncMock(return_value=mock_logs)
+ hook._async_conn = mock_client
+
+ 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"]
+
+ @pytest.mark.asyncio
+ async def test_get_logs_returns_none_sentinel_when_content_is_none(self,
async_conn_with_credentials):
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_logs = MagicMock(spec=Logs)
+ mock_logs.content = None
+ mock_client.containers.list_logs = AsyncMock(return_value=mock_logs)
+ hook._async_conn = mock_client
+
+ result = await hook.get_logs("my-rg", "my-container")
+ assert result == [None]
+
+ @pytest.mark.asyncio
+ async def test_delete(self, async_conn_with_credentials):
+ hook =
AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_client.container_groups.begin_delete = AsyncMock()
Review Comment:
We can change if we decide against `fire-and-forget`
--
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]
