def test_execute(self, aci_mock): aci_mock.return_value.get_state_exitcode_details.return_value = "Terminated", 0, "test" 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_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_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_with_failures(self, aci_mock): aci_mock.return_value.get_state_exitcode_details.return_value = "Terminated", 1, "test" 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_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_messages_logs(self, aci_mock): aci_mock.return_value.get_state_exitcode_details.side_effect = [("Running", 0, "test"), ("Terminated", 0, "test")] aci_mock.return_value.get_messages.return_value = ["test", "messages"] 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_exitcode_details.call_count, 2) self.assertEqual(aci_mock.return_value.get_messages.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_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) 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') self.assertEqual(called_cg.tags, tags) 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_messages_logs(self, aci_mock): aci_mock.return_value.get_state_exitcode_details.side_effect = [ ("Running", 0, "test"), ("Terminated", 0, "test") ] aci_mock.return_value.get_messages.return_value = ["test", "messages"] 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_exitcode_details.call_count, 2) self.assertEqual(aci_mock.return_value.get_messages.call_count, 2) self.assertEqual(aci_mock.return_value.get_logs.call_count, 2) self.assertEqual(aci_mock.return_value.delete.call_count, 1)
from airflow import DAG from airflow.contrib.operators.azure_container_instances_operator import AzureContainerInstancesOperator default_args = { 'owner': 'airflow', 'depends_on_past': False, 'start_date': datetime(2018, 11, 1), '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)) 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')