def test_execute_with_messages_logs(self, aci_mock):
events = [Event(message="test"), Event(message="messages")]
expected_c_state1 = ContainerState(state='Running', exit_code=0, detail_status='test')
expected_cg1 = make_mock_cg(expected_c_state1, events)
expected_c_state2 = ContainerState(state='Terminated', exit_code=0, detail_status='test')
expected_cg2 = make_mock_cg(expected_c_state2, events)
aci_mock.return_value.get_state.side_effect = [expected_cg1, expected_cg2]
aci_mock.return_value.get_logs.return_value = ["test", "logs"]
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',
)
aci.execute(None)
self.assertEqual(aci_mock.return_value.create_or_update.call_count, 1)
self.assertEqual(aci_mock.return_value.get_state.call_count, 2)
self.assertEqual(aci_mock.return_value.get_logs.call_count, 2)
self.assertEqual(aci_mock.return_value.delete.call_count, 1)
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)
aci_mock.return_value.get_state.return_value = expected_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',
)
aci.execute(None)
self.assertEqual(aci_mock.return_value.create_or_update.call_count, 1)
(called_rg, called_cn, called_cg), _ = aci_mock.return_value.create_or_update.call_args
self.assertEqual(called_rg, 'resource-group')
self.assertEqual(called_cn, 'container-name')
self.assertEqual(called_cg.location, 'region')
self.assertEqual(called_cg.image_registry_credentials, None)
self.assertEqual(called_cg.restart_policy, 'Never')
self.assertEqual(called_cg.os_type, 'Linux')
called_cg_container = called_cg.containers[0]
self.assertEqual(called_cg_container.name, 'container-name')
self.assertEqual(called_cg_container.image, 'container-image')
self.assertEqual(aci_mock.return_value.delete.call_count, 1)
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.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',
ip_address=ipaddress,
)
aci.execute(None)
self.assertEqual(aci_mock.return_value.create_or_update.call_count, 1)
(_, _, called_cg), _ = aci_mock.return_value.create_or_update.call_args
self.assertEqual(called_cg.ip_address, ipaddress)
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.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',
restart_policy="Always",
os_type='Windows',
)
aci.execute(None)
self.assertEqual(aci_mock.return_value.create_or_update.call_count, 1)
(_, _, called_cg), _ = aci_mock.return_value.create_or_update.call_args
self.assertEqual(called_cg.restart_policy, 'Always')
self.assertEqual(called_cg.os_type, 'Windows')
def test_name_checker(self):
valid_names = ['test-dash', 'name-with-length---63' * 3]
invalid_names = [
'test_underscore', 'name-with-length---84' * 4,
'name-ending-with-dash-', '-name-starting-with-dash'
]
for name in invalid_names:
with self.assertRaises(AirflowException):
AzureContainerInstancesOperator._check_name(name)
for name in valid_names:
checked_name = AzureContainerInstancesOperator._check_name(name)
self.assertEqual(checked_name, name)
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.exists.return_value = False
with self.assertRaises(AirflowException) as e:
AzureContainerInstancesOperator(
ci_conn_id=None,
registry_conn_id=None,
resource_group='resource-group',
name='container-name',
image='container-image',
region='region',
task_id='task',
restart_policy='Everyday',
)
self.assertEqual(
str(e.exception),
"Invalid value for the restart_policy argument. "
"Please set one of 'Always', 'OnFailure','Never' as the restart_policy. "
"Found `Everyday`",
)
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.exists.return_value = False
with self.assertRaises(AirflowException) as e:
AzureContainerInstancesOperator(
ci_conn_id=None,
registry_conn_id=None,
resource_group='resource-group',
name='container-name',
image='container-image',
region='region',
task_id='task',
os_type='MacOs',
)
self.assertEqual(
str(e.exception),
"Invalid value for the os_type argument. "
"Please set 'Linux' or 'Windows' as the os_type. "
"Found `MacOs`.",
)
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
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')
with self.assertRaises(AirflowException):
aci.execute(None)
self.assertEqual(aci_mock.return_value.delete.call_count, 1)
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)
tags = {"testKey": "testValue"}
aci_mock.return_value.get_state.return_value = expected_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',
tags=tags,
)
aci.execute(None)
assert aci_mock.return_value.create_or_update.call_count == 1
(called_rg, called_cn,
called_cg), _ = aci_mock.return_value.create_or_update.call_args
assert called_rg == 'resource-group'
assert called_cn == 'container-name'
assert called_cg.location == 'region'
assert called_cg.image_registry_credentials is None
assert called_cg.restart_policy == 'Never'
assert called_cg.os_type == 'Linux'
assert called_cg.tags == tags
called_cg_container = called_cg.containers[0]
assert called_cg_container.name == 'container-name'
assert called_cg_container.image == 'container-image'
assert aci_mock.return_value.delete.call_count == 1
'owner': 'airflow',
'depends_on_past': False,
'email': ['*****@*****.**'],
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': timedelta(minutes=5),
}
with DAG(
dag_id='aci_example',
default_args=default_args,
schedule_interval=timedelta(1),
start_date=datetime(2018, 11, 1),
tags=['example'],
) as dag:
t1 = AzureContainerInstancesOperator(
ci_conn_id='azure_container_instances_default',
registry_conn_id=None,
resource_group='resource-group',
name='aci-test-{{ ds }}',
image='hello-world',
region='WestUS2',
environment_variables={},
volumes=[],
memory_in_gb=4.0,
cpu=1.0,
task_id='start_container',
)
from airflow import DAG
from airflow.providers.microsoft.azure.operators.azure_container_instances import AzureContainerInstancesOperator
from datetime import datetime, timedelta
# Default settings applied to all tasks
default_args = {
'owner': 'airflow',
'depends_on_past': False,
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': timedelta(minutes=1)
}
with DAG('azure_container_instances',
start_date=datetime(2020, 12, 1),
max_active_runs=1,
schedule_interval='@daily',
default_args=default_args,
catchup=False) as dag:
opr_run_container = AzureContainerInstancesOperator(
task_id='run_container',
ci_conn_id='azure_container_conn_id',
registry_conn_id='azure_container_registry',
resource_group='CTO-ProtoLab-VideoAnalytics',
name='airflowcont',
image='flaskdemo.azurecr.io/flaskdemokunal:v1',
region='Central India',
fail_if_exists=False)
from airflow import DAG
from airflow.providers.microsoft.azure.operators.azure_container_instances import AzureContainerInstancesOperator
from datetime import datetime, timedelta
with DAG('azure_container_instances',
start_date=datetime(2020, 12, 1),
max_active_runs=1,
schedule_interval='@daily',
default_args={
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': timedelta(minutes=1)
},
catchup=False) as dag:
opr_run_container = AzureContainerInstancesOperator(
task_id='run_container',
ci_conn_id='azure_container_conn_id',
registry_conn_id=None,
resource_group='adf-tutorial',
name='astrotutorial',
image='hello-world:latest',
region='East US',
fail_if_exists=False)