def test_on_kill(self):
client_mock = mock.Mock(spec=APIClient)
operator = DockerSwarmOperator(image='',
auto_remove=False,
task_id='unittest')
operator.cli = client_mock
operator.service = {'ID': 'some_id'}
operator.on_kill()
client_mock.remove_service.assert_called_with('some_id')
import airflow
from airflow import DAG
from airflow.contrib.operators.docker_swarm_operator import DockerSwarmOperator
default_args = {
'owner': 'airflow',
'start_date': airflow.utils.dates.days_ago(1),
'email': ['*****@*****.**'],
'email_on_failure': True,
'email_on_retry': False
}
dag = DAG(
'docker_swarm_sample', # DAG ID
default_args=default_args,
schedule_interval='15 * * * *', # At 15th minute of every hour
catchup=False)
with dag as dag:
t1 = DockerSwarmOperator(
api_version='auto', # Docker API version
# command='/bin/sleep 45', # Command you want to run in the container
image='centos:latest', # The base image to use for running the container
auto_remove=
False, # Cleanup the container (and Docker service) once completed
task_id='swarm_start_container' # Unique task ID required by Airflow
)
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'}}]
yield [{'Status': {'State': 'complete'}}]
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.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)
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='******',
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')
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, 3)
client_mock.remove_service.assert_called_once_with('some_id')
default_args = {
'owner': 'airflow',
'depends_on_past': False,
'start_date': dt.datetime(2015, 6, 1),
'email': ['*****@*****.**'],
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': dt.timedelta(minutes=5),
# 'queue': 'bash_queue',
# 'pool': 'backfill',
# 'priority_weight': 10,
# 'end_date': datetime(2016, 1, 1),
}
with DAG('example_execute_docker_deploy_swarn',
catchup=False,
default_args=default_args,
schedule_interval=dt.timedelta(minutes=5)) as dag:
docker = DockerSwarmOperator(
api_version='auto',
docker_url='unix://var/run/docker.sock',
command='/bin/sleep 60',
image='centos:latest',
auto_remove=True,
task_id='sleep_with_swarm',
)
docker