Python DockerOperator.execute方法代码示例

# 需要导入模块: from airflow.operators.docker_operator import DockerOperator [as 别名]
# 或者: from airflow.operators.docker_operator.DockerOperator import execute [as 别名]
    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.logs.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'
        )

# 需要导入模块: from airflow.operators.docker_operator import DockerOperator [as 别名]
# 或者: from airflow.operators.docker_operator.DockerOperator import execute [as 别名]
    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.logs.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)

# 需要导入模块: from airflow.operators.docker_operator import DockerOperator [as 别名]
# 或者: from airflow.operators.docker_operator.DockerOperator import execute [as 别名]
    def test_execute(self, client_class_mock, mkdtemp_mock):
        host_config = mock.Mock()
        mkdtemp_mock.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.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)
        operator.execute(None)

        client_class_mock.assert_called_with(base_url='unix://var/run/docker.sock', tls=None,
                                             version='1.19')

        client_mock.create_container.assert_called_with(command='env',
                                                        environment={
                                                            'AIRFLOW_TMP_DIR': '/tmp/airflow',
                                                            'UNIT': 'TEST'
                                                        },
                                                        host_config=host_config,
                                                        image='ubuntu:latest',
                                                        user=None,
                                                        working_dir='/container/path'
                                                        )
        client_mock.create_host_config.assert_called_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)
        client_mock.images.assert_called_with(name='ubuntu:latest')
        client_mock.logs.assert_called_with(container='some_id', stream=True)
        client_mock.pull.assert_called_with('ubuntu:latest', stream=True)
        client_mock.wait.assert_called_with('some_id')

# 需要导入模块: from airflow.operators.docker_operator import DockerOperator [as 别名]
# 或者: from airflow.operators.docker_operator.DockerOperator import execute [as 别名]
    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.logs.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
        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()

# 需要导入模块: from airflow.operators.docker_operator import DockerOperator [as 别名]
# 或者: from airflow.operators.docker_operator.DockerOperator import execute [as 别名]
    def test_execute_with_docker_conn_id_use_hook(self, operator_client_mock):
        # Mock out a Docker client, so operations don't raise errors
        client_mock = mock.Mock(name='DockerOperator.Client mock', spec=Client)
        client_mock.images.return_value = []
        client_mock.create_container.return_value = {'Id': 'some_id'}
        client_mock.logs.return_value = []
        client_mock.pull.return_value = []
        client_mock.wait.return_value = 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.get_hook = mock.Mock(
            name='DockerOperator.get_hook mock',
            spec=DockerOperator.get_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.get_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'
        )

# 需要导入模块: from airflow.operators.docker_operator import DockerOperator [as 别名]
# 或者: from airflow.operators.docker_operator.DockerOperator import execute [as 别名]
    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.logs.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_with(assert_hostname=None, ca_cert='ca.pem',
                                          client_cert=('cert.pem', 'key.pem'),
                                          ssl_version=None, verify=True)

        client_class_mock.assert_called_with(base_url='https://127.0.0.1:2376',
                                             tls=tls_mock, version=None)