示例#1
0
    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)
示例#2
0
    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`',
        )
示例#3
0
    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',
    )
示例#5
0
    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",