diff --git a/airflow/providers/docker/operators/docker.py b/airflow/providers/docker/operators/docker.py index 0440d21b30c04..6bdeddf44ad03 100644 --- a/airflow/providers/docker/operators/docker.py +++ b/airflow/providers/docker/operators/docker.py @@ -219,7 +219,7 @@ def get_hook(self) -> DockerHook: tls=self.__get_tls_config(), ) - def _run_image(self) -> Optional[str]: + def _run_image(self) -> Optional[bytes]: """Run a Docker container with the provided image""" self.log.info('Starting docker container from image %s', self.image) @@ -269,14 +269,17 @@ def _run_image(self) -> Optional[str]: # duplicated conditional logic because of expensive operation ret = None if self.do_xcom_push: - ret = self.cli.logs(container=self.container['Id']) if self.xcom_all else line.encode('utf-8') + if self.xcom_all: + ret = self.cli.logs(container=self.container['Id']) + else: + ret = self.cli.logs(container=self.container['Id'], tail=1).strip() if self.auto_remove: self.cli.remove_container(self.container['Id']) return ret - def execute(self, context) -> Optional[str]: + def execute(self, context) -> Optional[bytes]: self.cli = self._get_cli() if not self.cli: raise Exception("The 'cli' should be initialized before!") diff --git a/tests/providers/docker/operators/test_docker.py b/tests/providers/docker/operators/test_docker.py index 0a2f8383e825f..e5ec762ac7bad 100644 --- a/tests/providers/docker/operators/test_docker.py +++ b/tests/providers/docker/operators/test_docker.py @@ -41,8 +41,8 @@ def setUp(self): self.client_mock = mock.Mock(spec=APIClient) self.client_mock.create_container.return_value = {'Id': 'some_id'} self.client_mock.images.return_value = [] - self.client_mock.attach.return_value = ['container log'] - self.client_mock.logs.return_value = ['container log'] + self.client_mock.attach.return_value = [b'container log'] + self.client_mock.logs.return_value = b'container log' # logs(..., stream=False) returns bytes self.client_mock.pull.return_value = {"status": "pull log"} self.client_mock.wait.return_value = {"StatusCode": 0} self.client_mock.create_host_config.return_value = mock.Mock()