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') 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_non_complete_service_raises_error(self, status, 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': status}}] 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) msg = "Service did not complete: {'ID': 'some_id'}" with pytest.raises(AirflowException) as ctx: operator.execute(None) assert str(ctx.value) == msg
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'}}] client_mock.configs.return_value = [] client_mock.secrets.return_value = [] 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) msg = "Service failed: {'ID': 'some_id'}" with self.assertRaises(AirflowException) as error: operator.execute(None) self.assertEqual(str(error.exception), msg)
def test_logging_with_requests_timeout(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.' raise requests.exceptions.ConnectionError('') 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, enable_logging=True, ) operator.execute(None) client_mock.service_logs.assert_called_once_with('some_id', follow=True, stdout=True, stderr=True, is_tty=True)
def test_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=True, task_id='unittest', enable_logging=False) operator.execute(None) client_mock.remove_service.assert_called_once_with('some_id')
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')