def test_failed_service_raises_error(self, types_mock, client_class_mock): mock_obj = mock.Mock() client_mock = mock.Mock(spec=APIClient) client_mock.create_service.return_value = {'ID': 'some_id'} client_mock.images.return_value = [] client_mock.pull.return_value = [b'{"status":"pull log"}'] client_mock.tasks.return_value = [{'Status': {'State': 'failed'}}] types_mock.TaskTemplate.return_value = mock_obj types_mock.ContainerSpec.return_value = mock_obj types_mock.RestartPolicy.return_value = mock_obj types_mock.Resources.return_value = mock_obj client_class_mock.return_value = client_mock operator = DockerSwarmOperator(image='', auto_remove=False, task_id='unittest') msg = "Service failed: {'ID': 'some_id'}" with self.assertRaises(AirflowException) as error: operator.execute(None) self.assertEqual(str(error.exception), msg)
def test_no_auto_remove(self, types_mock, client_class_mock): mock_obj = mock.Mock() client_mock = mock.Mock(spec=APIClient) client_mock.create_service.return_value = {'ID': 'some_id'} client_mock.images.return_value = [] client_mock.pull.return_value = [b'{"status":"pull log"}'] client_mock.tasks.return_value = [{'Status': {'State': 'complete'}}] types_mock.TaskTemplate.return_value = mock_obj types_mock.ContainerSpec.return_value = mock_obj types_mock.RestartPolicy.return_value = mock_obj types_mock.Resources.return_value = mock_obj client_class_mock.return_value = client_mock operator = DockerSwarmOperator(image='', auto_remove=False, task_id='unittest', enable_logging=False) operator.execute(None) self.assertEqual( client_mock.remove_service.call_count, 0, 'Docker service being removed even when `auto_remove` set to `False`', )
def test_on_kill(self): client_mock = mock.Mock(spec=APIClient) operator = DockerSwarmOperator(image='', auto_remove=False, task_id='unittest', enable_logging=False) operator.cli = client_mock operator.service = {'ID': 'some_id'} operator.on_kill() client_mock.remove_service.assert_called_once_with('some_id')
# under the License. from datetime import timedelta from airflow import DAG from airflow.providers.docker.operators.docker_swarm import DockerSwarmOperator from airflow.utils.dates import days_ago dag = DAG( 'docker_swarm_sample', default_args={ 'owner': 'airflow', 'depends_on_past': False, 'email': ['*****@*****.**'], 'email_on_failure': False, 'email_on_retry': False, }, schedule_interval=timedelta(minutes=10), start_date=days_ago(1), catchup=False, ) with dag as dag: t1 = DockerSwarmOperator( api_version='auto', docker_url='tcp://localhost:2375', # Set your docker URL command='/bin/sleep 10', image='centos:latest', auto_remove=True, task_id='sleep_with_swarm', )
def test_execute(self, types_mock, client_class_mock): mock_obj = mock.Mock() def _client_tasks_side_effect(): for _ in range(2): yield [{'Status': {'State': 'pending'}}] while True: yield [{'Status': {'State': 'complete'}}] def _client_service_logs_effect(): yield b'Testing is awesome.' client_mock = mock.Mock(spec=APIClient) client_mock.create_service.return_value = {'ID': 'some_id'} client_mock.service_logs.return_value = _client_service_logs_effect() client_mock.images.return_value = [] client_mock.pull.return_value = [b'{"status":"pull log"}'] client_mock.tasks.side_effect = _client_tasks_side_effect() types_mock.TaskTemplate.return_value = mock_obj types_mock.ContainerSpec.return_value = mock_obj types_mock.RestartPolicy.return_value = mock_obj types_mock.Resources.return_value = mock_obj client_class_mock.return_value = client_mock operator = DockerSwarmOperator( api_version='1.19', command='env', environment={'UNIT': 'TEST'}, image='ubuntu:latest', mem_limit='128m', user='******', task_id='unittest', auto_remove=True, tty=True, ) operator.execute(None) types_mock.TaskTemplate.assert_called_once_with( container_spec=mock_obj, restart_policy=mock_obj, resources=mock_obj ) types_mock.ContainerSpec.assert_called_once_with( image='ubuntu:latest', command='env', user='******', tty=True, env={'UNIT': 'TEST', 'AIRFLOW_TMP_DIR': '/tmp/airflow'}, ) types_mock.RestartPolicy.assert_called_once_with(condition='none') types_mock.Resources.assert_called_once_with(mem_limit='128m') client_class_mock.assert_called_once_with( base_url='unix://var/run/docker.sock', tls=None, version='1.19' ) client_mock.service_logs.assert_called_once_with( 'some_id', follow=True, stdout=True, stderr=True, is_tty=True ) csargs, cskwargs = client_mock.create_service.call_args_list[0] self.assertEqual( len(csargs), 1, 'create_service called with different number of arguments than expected' ) self.assertEqual(csargs, (mock_obj,)) self.assertEqual(cskwargs['labels'], {'name': 'airflow__adhoc_airflow__unittest'}) self.assertTrue(cskwargs['name'].startswith('airflow-')) self.assertEqual(client_mock.tasks.call_count, 5) client_mock.remove_service.assert_called_once_with('some_id')
'start_date': airflow.utils.dates.days_ago(1), 'email': ['*****@*****.**'], 'email_on_failure': False, 'email_on_retry': False } dag = DAG( 'docker_swarm_example2', # DAG ID default_args=default_args, schedule_interval='1 * * * *', # Run every hour catchup=False) with dag as dag: t1 = DockerSwarmOperator(task_id='swarm_sleep_30s', image='alpine:3.8', api_version='auto', command="/bin/sleep 10", auto_remove=True, do_xcom_push=False, enable_logging=False, tty=True) t2 = DockerSwarmOperator(task_id='swarm_sleep_15s', image='alpine:3.8', api_version='auto', command="/bin/sleep 15", auto_remove=True, do_xcom_push=False, enable_logging=False, tty=True) t3 = DockerSwarmOperator(task_id='swarm_sleep_20s', image='alpine:3.8', api_version='auto', command="/bin/sleep 10",