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",