def test_execute_with_docker_conn_id_use_hook(self, operator_client_mock,
operator_docker_hook):
# Mock out a Docker client, so operations don't raise errors
client_mock = mock.Mock(name='DockerOperator.APIClient mock',
spec=APIClient)
client_mock.images.return_value = []
client_mock.create_container.return_value = {'Id': 'some_id'}
client_mock.attach.return_value = []
client_mock.pull.return_value = []
client_mock.wait.return_value = {"StatusCode": 0}
operator_client_mock.return_value = client_mock
# Create the DockerOperator
operator = DockerOperator(image='publicregistry/someimage',
owner='unittest',
task_id='unittest',
docker_conn_id='some_conn_id')
# Mock out the DockerHook
hook_mock = mock.Mock(name='DockerHook mock', spec=DockerHook)
hook_mock.get_conn.return_value = client_mock
operator_docker_hook.return_value = hook_mock
operator.execute(None)
self.assertEqual(
operator_client_mock.call_count, 0,
'Client was called on the operator instead of the hook')
self.assertEqual(
operator_docker_hook.call_count, 1,
'Hook was not called although docker_conn_id configured')
self.assertEqual(client_mock.pull.call_count, 1,
'Image was not pulled using operator client')
def test_execute_no_docker_conn_id_no_hook(self, operator_client_mock):
# Mock out a Docker client, so operations don't raise errors
client_mock = mock.Mock(name='DockerOperator.APIClient mock',
spec=APIClient)
client_mock.images.return_value = []
client_mock.create_container.return_value = {'Id': 'some_id'}
client_mock.attach.return_value = []
client_mock.pull.return_value = []
client_mock.wait.return_value = {"StatusCode": 0}
operator_client_mock.return_value = client_mock
# Create the DockerOperator
operator = DockerOperator(image='publicregistry/someimage',
owner='unittest',
task_id='unittest')
# Mock out the DockerHook
hook_mock = mock.Mock(name='DockerHook mock', spec=DockerHook)
hook_mock.get_conn.return_value = client_mock
operator.get_hook = mock.Mock(name='DockerOperator.get_hook mock',
spec=DockerOperator.get_hook,
return_value=hook_mock)
operator.execute(None)
self.assertEqual(operator.get_hook.call_count, 0,
'Hook called though no docker_conn_id configured')
def test_execute_container_fails(self):
self.client_mock.wait.return_value = {"StatusCode": 1}
operator = DockerOperator(image='ubuntu',
owner='unittest',
task_id='unittest')
with self.assertRaises(AirflowException):
operator.execute(None)
def test_execute_tls(self, tls_class_mock):
tls_mock = mock.Mock()
tls_class_mock.return_value = tls_mock
operator = DockerOperator(
docker_url='tcp://127.0.0.1:2376',
image='ubuntu',
owner='unittest',
task_id='unittest',
tls_client_cert='cert.pem',
tls_ca_cert='ca.pem',
tls_client_key='key.pem',
)
operator.execute(None)
tls_class_mock.assert_called_once_with(
assert_hostname=None,
ca_cert='ca.pem',
client_cert=('cert.pem', 'key.pem'),
ssl_version=None,
verify=True,
)
self.client_class_mock.assert_called_once_with(
base_url='https://127.0.0.1:2376', tls=tls_mock, version=None)
def test_execute_tls(self, client_class_mock, tls_class_mock):
client_mock = mock.Mock(spec=APIClient)
client_mock.create_container.return_value = {'Id': 'some_id'}
client_mock.create_host_config.return_value = mock.Mock()
client_mock.images.return_value = []
client_mock.attach.return_value = []
client_mock.pull.return_value = []
client_mock.wait.return_value = {"StatusCode": 0}
client_class_mock.return_value = client_mock
tls_mock = mock.Mock()
tls_class_mock.return_value = tls_mock
operator = DockerOperator(docker_url='tcp://127.0.0.1:2376',
image='ubuntu',
owner='unittest',
task_id='unittest',
tls_client_cert='cert.pem',
tls_ca_cert='ca.pem',
tls_client_key='key.pem')
operator.execute(None)
tls_class_mock.assert_called_once_with(assert_hostname=None,
ca_cert='ca.pem',
client_cert=('cert.pem',
'key.pem'),
ssl_version=None,
verify=True)
client_class_mock.assert_called_once_with(
base_url='https://127.0.0.1:2376', tls=tls_mock, version=None)
def test_execute_xcom_behavior(self):
self.client_mock.pull.return_value = [b'{"status":"pull log"}']
kwargs = {
'api_version': '1.19',
'command': 'env',
'environment': {
'UNIT': 'TEST'
},
'private_environment': {
'PRIVATE': 'MESSAGE'
},
'image': 'ubuntu:latest',
'network_mode': 'bridge',
'owner': 'unittest',
'task_id': 'unittest',
'volumes': ['/host/path:/container/path'],
'working_dir': '/container/path',
'shm_size': 1000,
'host_tmp_dir': '/host/airflow',
'container_name': 'test_container',
'tty': True,
}
xcom_push_operator = DockerOperator(**kwargs, do_xcom_push=True)
no_xcom_push_operator = DockerOperator(**kwargs, do_xcom_push=False)
xcom_push_result = xcom_push_operator.execute(None)
no_xcom_push_result = no_xcom_push_operator.execute(None)
self.assertEqual(xcom_push_result, b'container log')
self.assertIs(no_xcom_push_result, None)
def test_auto_remove_container_fails(self):
self.client_mock.wait.return_value = {"StatusCode": 1}
operator = DockerOperator(image='ubuntu', owner='unittest', task_id='unittest', auto_remove=True)
operator.container = {'Id': 'some_id'}
with pytest.raises(AirflowException):
operator.execute(None)
self.client_mock.remove_container.assert_called_once_with('some_id')
def test_privileged(self):
privileged = mock.Mock()
operator = DockerOperator(task_id='test', image='test', privileged=privileged)
operator.execute(None)
self.client_mock.create_container.assert_called_once()
assert 'host_config' in self.client_mock.create_container.call_args[1]
assert 'privileged' in self.client_mock.create_host_config.call_args[1]
assert privileged is self.client_mock.create_host_config.call_args[1]['privileged']
def test_extra_hosts(self):
hosts_obj = mock.Mock()
operator = DockerOperator(task_id='test', image='test', extra_hosts=hosts_obj)
operator.execute(None)
self.client_mock.create_container.assert_called_once()
assert 'host_config' in self.client_mock.create_container.call_args[1]
assert 'extra_hosts' in self.client_mock.create_host_config.call_args[1]
assert hosts_obj is self.client_mock.create_host_config.call_args[1]['extra_hosts']
def test_execute(self):
operator = DockerOperator(
api_version='1.19',
command='env',
environment={'UNIT': 'TEST'},
private_environment={'PRIVATE': 'MESSAGE'},
image='ubuntu:latest',
network_mode='bridge',
owner='unittest',
task_id='unittest',
volumes=['/host/path:/container/path'],
entrypoint='["sh", "-c"]',
working_dir='/container/path',
shm_size=1000,
host_tmp_dir='/host/airflow',
container_name='test_container',
tty=True,
)
operator.execute(None)
self.client_class_mock.assert_called_once_with(
base_url='unix://var/run/docker.sock', tls=None, version='1.19'
)
self.client_mock.create_container.assert_called_once_with(
command='env',
name='test_container',
environment={'AIRFLOW_TMP_DIR': '/tmp/airflow', 'UNIT': 'TEST', 'PRIVATE': 'MESSAGE'},
host_config=self.client_mock.create_host_config.return_value,
image='ubuntu:latest',
user=None,
entrypoint=['sh', '-c'],
working_dir='/container/path',
tty=True,
)
self.client_mock.create_host_config.assert_called_once_with(
binds=['/host/path:/container/path', '/mkdtemp:/tmp/airflow'],
network_mode='bridge',
shm_size=1000,
cpu_shares=1024,
mem_limit=None,
auto_remove=False,
dns=None,
dns_search=None,
cap_add=None,
extra_hosts=None,
privileged=False,
)
self.tempdir_mock.assert_called_once_with(dir='/host/airflow', prefix='airflowtmp')
self.client_mock.images.assert_called_once_with(name='ubuntu:latest')
self.client_mock.attach.assert_called_once_with(
container='some_id', stdout=True, stderr=True, stream=True
)
self.client_mock.pull.assert_called_once_with('ubuntu:latest', stream=True, decode=True)
self.client_mock.wait.assert_called_once_with('some_id')
assert (
operator.cli.pull('ubuntu:latest', stream=True, decode=True) == self.client_mock.pull.return_value
)
def test_execute(self, client_class_mock, tempdir_mock):
host_config = mock.Mock()
tempdir_mock.return_value.__enter__.return_value = '/mkdtemp'
client_mock = mock.Mock(spec=APIClient)
client_mock.create_container.return_value = {'Id': 'some_id'}
client_mock.create_host_config.return_value = host_config
client_mock.images.return_value = []
client_mock.attach.return_value = ['container log']
client_mock.logs.return_value = ['container log']
client_mock.pull.return_value = [b'{"status":"pull log"}']
client_mock.wait.return_value = {"StatusCode": 0}
client_class_mock.return_value = client_mock
operator = DockerOperator(api_version='1.19', command='env', environment={'UNIT': 'TEST'},
image='ubuntu:latest', network_mode='bridge', owner='unittest',
task_id='unittest', volumes=['/host/path:/container/path'],
working_dir='/container/path', shm_size=1000,
host_tmp_dir='/host/airflow', container_name='test_container',
tty=True)
operator.execute(None)
client_class_mock.assert_called_once_with(base_url='unix://var/run/docker.sock', tls=None,
version='1.19')
client_mock.create_container.assert_called_once_with(command='env',
name='test_container',
environment={
'AIRFLOW_TMP_DIR': '/tmp/airflow',
'UNIT': 'TEST'
},
host_config=host_config,
image='ubuntu:latest',
user=None,
working_dir='/container/path',
tty=True
)
client_mock.create_host_config.assert_called_once_with(binds=['/host/path:/container/path',
'/mkdtemp:/tmp/airflow'],
network_mode='bridge',
shm_size=1000,
cpu_shares=1024,
mem_limit=None,
auto_remove=False,
dns=None,
dns_search=None)
tempdir_mock.assert_called_once_with(dir='/host/airflow', prefix='airflowtmp')
client_mock.images.assert_called_once_with(name='ubuntu:latest')
client_mock.attach.assert_called_once_with(container='some_id', stdout=True,
stderr=True, stream=True)
client_mock.pull.assert_called_once_with('ubuntu:latest', stream=True,
decode=True)
client_mock.wait.assert_called_once_with('some_id')
def test_execute_unicode_logs(self):
self.client_mock.attach.return_value = ['unicode container log 😁']
originalRaiseExceptions = logging.raiseExceptions # pylint: disable=invalid-name
logging.raiseExceptions = True
operator = DockerOperator(image='ubuntu', owner='unittest', task_id='unittest')
with mock.patch('traceback.print_exception') as print_exception_mock:
operator.execute(None)
logging.raiseExceptions = originalRaiseExceptions
print_exception_mock.assert_not_called()
def test_execute_no_docker_conn_id_no_hook(self):
# Create the DockerOperator
operator = DockerOperator(image='publicregistry/someimage', owner='unittest', task_id='unittest')
# Mock out the DockerHook
hook_mock = mock.Mock(name='DockerHook mock', spec=DockerHook)
hook_mock.get_conn.return_value = self.client_mock
operator.get_hook = mock.Mock(
name='DockerOperator.get_hook mock', spec=DockerOperator.get_hook, return_value=hook_mock
)
operator.execute(None)
assert operator.get_hook.call_count == 0, 'Hook called though no docker_conn_id configured'
def test_execute_container_fails(self, client_class_mock):
client_mock = mock.Mock(spec=APIClient)
client_mock.create_container.return_value = {'Id': 'some_id'}
client_mock.create_host_config.return_value = mock.Mock()
client_mock.images.return_value = []
client_mock.attach.return_value = []
client_mock.pull.return_value = []
client_mock.wait.return_value = {"StatusCode": 1}
client_class_mock.return_value = client_mock
operator = DockerOperator(image='ubuntu', owner='unittest', task_id='unittest')
with self.assertRaises(AirflowException):
operator.execute(None)
def test_execute_with_docker_conn_id_use_hook(self, hook_class_mock):
# Create the DockerOperator
operator = DockerOperator(
image='publicregistry/someimage',
owner='unittest',
task_id='unittest',
docker_conn_id='some_conn_id',
)
# Mock out the DockerHook
hook_mock = mock.Mock(name='DockerHook mock', spec=DockerHook)
hook_mock.get_conn.return_value = self.client_mock
hook_class_mock.return_value = hook_mock
operator.execute(None)
assert self.client_class_mock.call_count == 0, 'Client was called on the operator instead of the hook'
assert hook_class_mock.call_count == 1, 'Hook was not called although docker_conn_id configured'
assert self.client_mock.pull.call_count == 1, 'Image was not pulled using operator client'
def test_extra_hosts(self):
hosts_obj = mock.Mock()
operator = DockerOperator(task_id='test',
image='test',
extra_hosts=hosts_obj)
operator.execute(None)
self.client_mock.create_container.assert_called_once()
self.assertIn(
'host_config',
self.client_mock.create_container.call_args.kwargs,
)
self.assertIn(
'extra_hosts',
self.client_mock.create_host_config.call_args.kwargs,
)
self.assertIs(
hosts_obj,
self.client_mock.create_host_config.call_args.
kwargs['extra_hosts'],
)
def test_execute_unicode_logs(self, client_class_mock):
client_mock = mock.Mock(spec=APIClient)
client_mock.create_container.return_value = {'Id': 'some_id'}
client_mock.create_host_config.return_value = mock.Mock()
client_mock.images.return_value = []
client_mock.attach.return_value = ['unicode container log 😁']
client_mock.pull.return_value = []
client_mock.wait.return_value = {"StatusCode": 0}
client_class_mock.return_value = client_mock
originalRaiseExceptions = logging.raiseExceptions # pylint: disable=invalid-name
logging.raiseExceptions = True
operator = DockerOperator(image='ubuntu', owner='unittest', task_id='unittest')
with mock.patch('traceback.print_exception') as print_exception_mock:
operator.execute(None)
logging.raiseExceptions = originalRaiseExceptions
print_exception_mock.assert_not_called()
def test_execute_xcom_behavior(self, client_class_mock, tempdir_mock):
tempdir_mock.return_value.__enter__.return_value = '/mkdtemp'
client_mock = mock.Mock(spec=APIClient)
client_mock.images.return_value = []
client_mock.create_container.return_value = {'Id': 'some_id'}
client_mock.attach.return_value = ['container log']
client_mock.pull.return_value = [b'{"status":"pull log"}']
client_mock.wait.return_value = {"StatusCode": 0}
client_class_mock.return_value = client_mock
kwargs = {
'api_version': '1.19',
'command': 'env',
'environment': {
'UNIT': 'TEST'
},
'private_environment': {
'PRIVATE': 'MESSAGE'
},
'image': 'ubuntu:latest',
'network_mode': 'bridge',
'owner': 'unittest',
'task_id': 'unittest',
'volumes': ['/host/path:/container/path'],
'working_dir': '/container/path',
'shm_size': 1000,
'host_tmp_dir': '/host/airflow',
'container_name': 'test_container',
'tty': True,
}
xcom_push_operator = DockerOperator(**kwargs, do_xcom_push=True)
no_xcom_push_operator = DockerOperator(**kwargs, do_xcom_push=False)
xcom_push_result = xcom_push_operator.execute(None)
no_xcom_push_result = no_xcom_push_operator.execute(None)
self.assertEqual(xcom_push_result, b'container log')
self.assertIs(no_xcom_push_result, None)
