From 95a5cb2d085a549d75015bf8a72ef22ce30b5845 Mon Sep 17 00:00:00 2001 From: rpkay Date: Fri, 17 Nov 2023 21:26:34 +0530 Subject: [PATCH] re-wrote test case based on feedback --- tests/providers/docker/test_exceptions.py | 66 +++++++++++++++++------ 1 file changed, 50 insertions(+), 16 deletions(-) diff --git a/tests/providers/docker/test_exceptions.py b/tests/providers/docker/test_exceptions.py index 4fbc4694c55c0..1f3653c39dcd8 100644 --- a/tests/providers/docker/test_exceptions.py +++ b/tests/providers/docker/test_exceptions.py @@ -16,31 +16,65 @@ # specific language governing permissions and limitations # under the License. """Test for Exceptions used by Docker provider.""" + from __future__ import annotations +from unittest import mock + import pytest +from docker import APIClient from airflow.providers.docker.exceptions import ( DockerContainerFailedException, DockerContainerFailedSkipException, ) +from airflow.providers.docker.operators.docker import DockerOperator -FAILED_MESSAGE = "Error 1" -FAILED_LOGS = ["Log 1", "Log2"] -FAILED_SKIP_MESSAGE = "Error 2" -FAILED_SKIP_LOGS = ["Log 3", "Log4"] +FAILED_MESSAGE = {"StatusCode": 1} +FAILED_LOGS = ["unicode container log 😁 ", b"byte string container log"] +EXPECTED_MESSAGE = f"Docker container failed: {FAILED_MESSAGE}" +FAILED_SKIP_MESSAGE = {"StatusCode": 2} +SKIP_ON_EXIT_CODE = 2 +EXPECTED_SKIP_MESSAGE = f"Docker container returned exit code {[SKIP_ON_EXIT_CODE]}. Skipping." +@pytest.mark.parametrize( + "failed_msg, log_line, expected_message, skip_on_exit_code", + [ + (FAILED_MESSAGE, FAILED_LOGS, EXPECTED_MESSAGE, None), + (FAILED_SKIP_MESSAGE, FAILED_LOGS, EXPECTED_SKIP_MESSAGE, SKIP_ON_EXIT_CODE), + ], +) class TestDockerContainerExceptions: - @pytest.mark.parametrize( - "exception_class, message, logs", - [ - (DockerContainerFailedException, FAILED_MESSAGE, FAILED_LOGS), - (DockerContainerFailedSkipException, FAILED_SKIP_MESSAGE, FAILED_SKIP_LOGS), - ], - ) - def test_docker_failed_exception(self, exception_class, message, logs): - with pytest.raises(exception_class) as exc_info: - raise exception_class(message, logs) - assert exc_info.value.args[0] == message - assert exc_info.value.logs == logs + @pytest.fixture(autouse=True, scope="function") + def setup_patchers(self, docker_api_client_patcher): + self.client_mock = mock.MagicMock(spec=APIClient) + self.client_mock.wait.return_value = {"StatusCode": 0} + self.log_messages = ["container log 😁 ", b"byte string container log"] + self.client_mock.attach.return_value = self.log_messages + + self.client_mock.logs.side_effect = ( + lambda **kwargs: iter(self.log_messages[-kwargs["tail"] :]) + if "tail" in kwargs + else iter(self.log_messages) + ) + + docker_api_client_patcher.return_value = self.client_mock + + def test_docker_failed_exception(self, failed_msg, log_line, expected_message, skip_on_exit_code): + self.client_mock.attach.return_value = log_line + self.client_mock.wait.return_value = failed_msg + + operator = DockerOperator( + image="ubuntu", owner="unittest", task_id="unittest", skip_on_exit_code=skip_on_exit_code + ) + + if skip_on_exit_code: + with pytest.raises(DockerContainerFailedSkipException) as raised_exception: + operator.execute(None) + else: + with pytest.raises(DockerContainerFailedException) as raised_exception: + operator.execute(None) + + assert str(raised_exception.value) == expected_message + assert raised_exception.value.logs == [log_line[0].strip(), log_line[1].decode("utf-8")]