airbyte-cdk 0.61.2__py3-none-any.whl → 0.62.1__py3-none-any.whl

unit_tests/sources/test_abstract_source.py

@@ -13,7 +13,6 @@ import pytest
 from airbyte_cdk.models import (
     AirbyteCatalog,
     AirbyteConnectionStatus,
-    AirbyteErrorTraceMessage,
     AirbyteLogMessage,
     AirbyteMessage,
     AirbyteRecordMessage,
@@ -28,7 +27,6 @@ from airbyte_cdk.models import (
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteStream,
     DestinationSyncMode,
-    FailureType,
     Level,
     Status,
     StreamDescriptor,
@@ -42,7 +40,6 @@ from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.streams import IncrementalMixin, Stream
 from airbyte_cdk.sources.utils.record_helper import stream_data_to_airbyte_message
-from airbyte_cdk.utils.airbyte_secrets_utils import update_secrets
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 from pytest import fixture
 
@@ -57,14 +54,12 @@ class MockSource(AbstractSource):
         per_stream: bool = True,
         message_repository: MessageRepository = None,
         exception_on_missing_stream: bool = True,
-        stop_sync_on_stream_failure: bool = False,
     ):
         self._streams = streams
         self.check_lambda = check_lambda
         self.per_stream = per_stream
         self.exception_on_missing_stream = exception_on_missing_stream
         self._message_repository = message_repository
-        self._stop_sync_on_stream_failure = stop_sync_on_stream_failure
 
     def check_connection(self, logger: logging.Logger, config: Mapping[str, Any]) -> Tuple[bool, Optional[Any]]:
         if self.check_lambda:
@@ -89,12 +84,6 @@ class MockSource(AbstractSource):
         return self._message_repository
 
 
-class MockSourceWithStopSyncFalseOverride(MockSource):
-    @property
-    def stop_sync_on_stream_failure(self) -> bool:
-        return False
-
-
 class StreamNoStateMethod(Stream):
     name = "managers"
     primary_key = None
@@ -126,11 +115,8 @@ class StreamRaisesException(Stream):
     name = "lamentations"
     primary_key = None
 
-    def __init__(self, exception_to_raise):
-        self._exception_to_raise = exception_to_raise
-
     def read_records(self, *args, **kwargs) -> Iterable[Mapping[str, Any]]:
-        raise self._exception_to_raise
+        raise AirbyteTracedException(message="I was born only to crash like Icarus")
 
 
 MESSAGE_FROM_REPOSITORY = Mock()
@@ -305,7 +291,7 @@ def test_read_stream_emits_repository_message_on_error(mocker, message_repositor
 
     source = MockSource(streams=[stream], message_repository=message_repository)
 
-    with pytest.raises(AirbyteTracedException):
+    with pytest.raises(RuntimeError):
         messages = list(source.read(logger, {}, ConfiguredAirbyteCatalog(streams=[_configured_stream(stream, SyncMode.full_refresh)])))
         assert MESSAGE_FROM_REPOSITORY in messages
 
@@ -320,14 +306,14 @@ def test_read_stream_with_error_gets_display_message(mocker):
     catalog = ConfiguredAirbyteCatalog(streams=[_configured_stream(stream, SyncMode.full_refresh)])
 
     # without get_error_display_message
-    with pytest.raises(AirbyteTracedException):
+    with pytest.raises(RuntimeError, match="oh no!"):
         list(source.read(logger, {}, catalog))
 
     mocker.patch.object(MockStream, "get_error_display_message", return_value="my message")
 
-    with pytest.raises(AirbyteTracedException) as exc:
+    with pytest.raises(AirbyteTracedException, match="oh no!") as exc:
         list(source.read(logger, {}, catalog))
-    assert "oh no!" in exc.value.message
+    assert exc.value.message == "my message"
 
 
 GLOBAL_EMITTED_AT = 1
@@ -372,22 +358,6 @@ def _as_state(state_data: Dict[str, Any], stream_name: str = "", per_stream_stat
     return AirbyteMessage(type=Type.STATE, state=AirbyteStateMessage(data=state_data))
 
 
-def _as_error_trace(stream: str, error_message: str,  internal_message: Optional[str], failure_type: Optional[FailureType], stack_trace: Optional[str]) -> AirbyteMessage:
-    trace_message = AirbyteTraceMessage(
-        emitted_at=datetime.datetime.now().timestamp() * 1000.0,
-        type=TraceType.ERROR,
-        error=AirbyteErrorTraceMessage(
-                stream_descriptor=StreamDescriptor(name=stream),
-                message=error_message,
-                internal_message=internal_message,
-                failure_type=failure_type,
-                stack_trace=stack_trace,
-            ),
-    )
-
-    return AirbyteMessage(type=MessageType.TRACE, trace=trace_message)
-
-
 def _configured_stream(stream: Stream, sync_mode: SyncMode):
     return ConfiguredAirbyteStream(
         stream=stream.as_airbyte_stream(),
@@ -1184,7 +1154,13 @@ class TestIncrementalRead:
 def test_checkpoint_state_from_stream_instance():
     teams_stream = MockStreamOverridesStateMethod()
     managers_stream = StreamNoStateMethod()
-    state_manager = ConnectorStateManager({"teams": teams_stream, "managers": managers_stream}, [])
+    state_manager = ConnectorStateManager(
+        {
+            "teams": AirbyteStream(name="teams", namespace="", json_schema={}, supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental]),
+            "managers": AirbyteStream(name="managers", namespace="", json_schema={}, supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental])
+        },
+        [],
+    )
 
     # The stream_state passed to checkpoint_state() should be ignored since stream implements state function
     teams_stream.state = {"updated_at": "2022-09-11"}
@@ -1198,124 +1174,21 @@ def test_checkpoint_state_from_stream_instance():
     )
 
 
-@pytest.mark.parametrize(
-    "exception_to_raise,expected_error_message,expected_internal_message",
-    [
-        pytest.param(AirbyteTracedException(message="I was born only to crash like Icarus"), "I was born only to crash like Icarus", None, id="test_raises_traced_exception"),
-        pytest.param(Exception("Generic connector error message"), "Something went wrong in the connector. See the logs for more details.", "Generic connector error message", id="test_raises_generic_exception"),
-    ]
-)
-def test_continue_sync_with_failed_streams(mocker, exception_to_raise, expected_error_message, expected_internal_message):
-    """
-    Tests that running a sync for a connector with multiple streams will continue syncing when one stream fails
-    with an error. This source does not override the default behavior defined in the AbstractSource class.
-    """
-    stream_output = [{"k1": "v1"}, {"k2": "v2"}]
-    s1 = MockStream([({"sync_mode": SyncMode.full_refresh}, stream_output)], name="s1")
-    s2 = StreamRaisesException(exception_to_raise=exception_to_raise)
-    s3 = MockStream([({"sync_mode": SyncMode.full_refresh}, stream_output)], name="s3")
-
-    mocker.patch.object(MockStream, "get_json_schema", return_value={})
-    mocker.patch.object(StreamRaisesException, "get_json_schema", return_value={})
-
-    src = MockSource(streams=[s1, s2, s3])
-    catalog = ConfiguredAirbyteCatalog(
-        streams=[
-            _configured_stream(s1, SyncMode.full_refresh),
-            _configured_stream(s2, SyncMode.full_refresh),
-            _configured_stream(s3, SyncMode.full_refresh),
-        ]
-    )
-
-    expected = _fix_emitted_at(
-        [
-            _as_stream_status("s1", AirbyteStreamStatus.STARTED),
-            _as_stream_status("s1", AirbyteStreamStatus.RUNNING),
-            *_as_records("s1", stream_output),
-            _as_stream_status("s1", AirbyteStreamStatus.COMPLETE),
-            _as_stream_status("lamentations", AirbyteStreamStatus.STARTED),
-            _as_stream_status("lamentations", AirbyteStreamStatus.INCOMPLETE),
-            _as_error_trace("lamentations", expected_error_message, expected_internal_message, FailureType.system_error, None),
-            _as_stream_status("s3", AirbyteStreamStatus.STARTED),
-            _as_stream_status("s3", AirbyteStreamStatus.RUNNING),
-            *_as_records("s3", stream_output),
-            _as_stream_status("s3", AirbyteStreamStatus.COMPLETE),
-        ]
-    )
-
-    with pytest.raises(AirbyteTracedException) as exc:
-        messages = [_remove_stack_trace(message) for message in src.read(logger, {}, catalog)]
-        messages = _fix_emitted_at(messages)
-
-        assert expected == messages
-
-    assert "lamentations" in exc.value.message
-
-
-def test_continue_sync_source_override_false(mocker):
-    """
-    Tests that running a sync for a connector explicitly overriding the default AbstractSource.stop_sync_on_stream_failure
-    property to be False which will continue syncing stream even if one encountered an exception.
-    """
-    update_secrets(["API_KEY_VALUE"])
-
-    stream_output = [{"k1": "v1"}, {"k2": "v2"}]
-    s1 = MockStream([({"sync_mode": SyncMode.full_refresh}, stream_output)], name="s1")
-    s2 = StreamRaisesException(exception_to_raise=AirbyteTracedException(message="I was born only to crash like Icarus"))
-    s3 = MockStream([({"sync_mode": SyncMode.full_refresh}, stream_output)], name="s3")
-
-    mocker.patch.object(MockStream, "get_json_schema", return_value={})
-    mocker.patch.object(StreamRaisesException, "get_json_schema", return_value={})
-
-    src = MockSourceWithStopSyncFalseOverride(streams=[s1, s2, s3])
-    catalog = ConfiguredAirbyteCatalog(
-        streams=[
-            _configured_stream(s1, SyncMode.full_refresh),
-            _configured_stream(s2, SyncMode.full_refresh),
-            _configured_stream(s3, SyncMode.full_refresh),
-        ]
-    )
-
-    expected = _fix_emitted_at(
-        [
-            _as_stream_status("s1", AirbyteStreamStatus.STARTED),
-            _as_stream_status("s1", AirbyteStreamStatus.RUNNING),
-            *_as_records("s1", stream_output),
-            _as_stream_status("s1", AirbyteStreamStatus.COMPLETE),
-            _as_stream_status("lamentations", AirbyteStreamStatus.STARTED),
-            _as_stream_status("lamentations", AirbyteStreamStatus.INCOMPLETE),
-            _as_error_trace("lamentations", "I was born only to crash like Icarus", None, FailureType.system_error, None),
-            _as_stream_status("s3", AirbyteStreamStatus.STARTED),
-            _as_stream_status("s3", AirbyteStreamStatus.RUNNING),
-            *_as_records("s3", stream_output),
-            _as_stream_status("s3", AirbyteStreamStatus.COMPLETE),
-        ]
-    )
-
-    with pytest.raises(AirbyteTracedException) as exc:
-        messages = [_remove_stack_trace(message) for message in src.read(logger, {}, catalog)]
-        messages = _fix_emitted_at(messages)
-
-        assert expected == messages
-
-    assert "lamentations" in exc.value.message
-
-
-def test_sync_error_trace_messages_obfuscate_secrets(mocker):
+def test_continue_sync_with_failed_streams(mocker):
     """
-    Tests that exceptions emitted as trace messages by a source have secrets properly sanitized
+    Tests that running a sync for a connector with multiple streams and continue_sync_on_stream_failure enabled continues
+    syncing even when one stream fails with an error.
     """
-    update_secrets(["API_KEY_VALUE"])
-
     stream_output = [{"k1": "v1"}, {"k2": "v2"}]
     s1 = MockStream([({"sync_mode": SyncMode.full_refresh}, stream_output)], name="s1")
-    s2 = StreamRaisesException(exception_to_raise=AirbyteTracedException(message="My api_key value API_KEY_VALUE flew too close to the sun."))
+    s2 = StreamRaisesException()
     s3 = MockStream([({"sync_mode": SyncMode.full_refresh}, stream_output)], name="s3")
 
     mocker.patch.object(MockStream, "get_json_schema", return_value={})
     mocker.patch.object(StreamRaisesException, "get_json_schema", return_value={})
 
     src = MockSource(streams=[s1, s2, s3])
+    mocker.patch.object(MockSource, "continue_sync_on_stream_failure", return_value=True)
     catalog = ConfiguredAirbyteCatalog(
         streams=[
             _configured_stream(s1, SyncMode.full_refresh),
@@ -1332,7 +1205,6 @@ def test_sync_error_trace_messages_obfuscate_secrets(mocker):
             _as_stream_status("s1", AirbyteStreamStatus.COMPLETE),
             _as_stream_status("lamentations", AirbyteStreamStatus.STARTED),
             _as_stream_status("lamentations", AirbyteStreamStatus.INCOMPLETE),
-            _as_error_trace("lamentations", "My api_key value **** flew too close to the sun.", None, FailureType.system_error, None),
             _as_stream_status("s3", AirbyteStreamStatus.STARTED),
             _as_stream_status("s3", AirbyteStreamStatus.RUNNING),
             *_as_records("s3", stream_output),
@@ -1340,30 +1212,32 @@ def test_sync_error_trace_messages_obfuscate_secrets(mocker):
         ]
     )
 
+    messages = []
     with pytest.raises(AirbyteTracedException) as exc:
-        messages = [_remove_stack_trace(message) for message in src.read(logger, {}, catalog)]
-        messages = _fix_emitted_at(messages)
-
-        assert expected == messages
+        # We can't use list comprehension or list() here because we are still raising a final exception for the
+        # failed streams and that disrupts parsing the generator into the messages emitted before
+        for message in src.read(logger, {}, catalog):
+            messages.append(message)
 
+    messages = _fix_emitted_at(messages)
+    assert expected == messages
     assert "lamentations" in exc.value.message
 
 
-def test_continue_sync_with_failed_streams_with_override_false(mocker):
+def test_stop_sync_with_failed_streams(mocker):
     """
-    Tests that running a sync for a connector with multiple streams and stop_sync_on_stream_failure enabled stops
-    the sync when one stream fails with an error.
+    Tests that running a sync for a connector with multiple streams and continue_sync_on_stream_failure disabled stops
+    syncing once a stream fails with an error.
     """
     stream_output = [{"k1": "v1"}, {"k2": "v2"}]
     s1 = MockStream([({"sync_mode": SyncMode.full_refresh}, stream_output)], name="s1")
-    s2 = StreamRaisesException(AirbyteTracedException(message="I was born only to crash like Icarus"))
+    s2 = StreamRaisesException()
     s3 = MockStream([({"sync_mode": SyncMode.full_refresh}, stream_output)], name="s3")
 
     mocker.patch.object(MockStream, "get_json_schema", return_value={})
     mocker.patch.object(StreamRaisesException, "get_json_schema", return_value={})
 
     src = MockSource(streams=[s1, s2, s3])
-    mocker.patch.object(MockSource, "stop_sync_on_stream_failure", return_value=True)
     catalog = ConfiguredAirbyteCatalog(
         streams=[
             _configured_stream(s1, SyncMode.full_refresh),
@@ -1380,23 +1254,15 @@ def test_continue_sync_with_failed_streams_with_override_false(mocker):
             _as_stream_status("s1", AirbyteStreamStatus.COMPLETE),
             _as_stream_status("lamentations", AirbyteStreamStatus.STARTED),
             _as_stream_status("lamentations", AirbyteStreamStatus.INCOMPLETE),
-            _as_error_trace("lamentations", "I was born only to crash like Icarus", None, FailureType.system_error, None),
         ]
     )
 
-    with pytest.raises(AirbyteTracedException) as exc:
-        messages = [_remove_stack_trace(message) for message in src.read(logger, {}, catalog)]
-        messages = _fix_emitted_at(messages)
-
-        assert expected == messages
-
-    assert "lamentations" in exc.value.message
-
+    messages = []
+    with pytest.raises(AirbyteTracedException):
+        # We can't use list comprehension or list() here because we are still raising a final exception for the
+        # failed streams and that disrupts parsing the generator into the messages emitted before
+        for message in src.read(logger, {}, catalog):
+            messages.append(message)
 
-def _remove_stack_trace(message: AirbyteMessage) -> AirbyteMessage:
-    """
-    Helper method that removes the stack trace from Airbyte trace messages to make asserting against expected records easier
-    """
-    if message.trace and message.trace.error and message.trace.error.stack_trace:
-        message.trace.error.stack_trace = None
-    return message
+    messages = _fix_emitted_at(messages)
+    assert expected == messages

unit_tests/sources/test_connector_state_manager.py

@@ -3,21 +3,21 @@
 #
 
 from contextlib import nullcontext as does_not_raise
-from typing import Any, Iterable, List, Mapping
+from typing import List
 
 import pytest
-from airbyte_cdk.models import AirbyteMessage, AirbyteStateBlob, AirbyteStateMessage, AirbyteStateType, AirbyteStreamState, StreamDescriptor
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    AirbyteStateBlob,
+    AirbyteStateMessage,
+    AirbyteStateType,
+    AirbyteStream,
+    AirbyteStreamState,
+    StreamDescriptor,
+    SyncMode,
+)
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager, HashableStreamDescriptor
-from airbyte_cdk.sources.streams import Stream
-
-
-class StreamWithNamespace(Stream):
-    primary_key = None
-    namespace = "public"
-
-    def read_records(self, *args, **kwargs) -> Iterable[Mapping[str, Any]]:
-        return {}
 
 
 @pytest.mark.parametrize(
@@ -156,7 +156,11 @@ class StreamWithNamespace(Stream):
     ),
 )
 def test_initialize_state_manager(input_stream_state, expected_stream_state, expected_error):
-    stream_to_instance_map = {"actors": StreamWithNamespace()}
+    stream_to_instance_map = {
+        "actors": AirbyteStream(
+            name="actors", namespace="public", json_schema={}, supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental],
+        )
+    }
 
     if isinstance(input_stream_state, List):
         input_stream_state = [AirbyteStateMessage.parse_obj(state_obj) for state_obj in list(input_stream_state)]
@@ -264,7 +268,10 @@ def test_initialize_state_manager(input_stream_state, expected_stream_state, exp
     ],
 )
 def test_get_stream_state(input_state, stream_name, namespace, expected_state):
-    stream_to_instance_map = {"users": StreamWithNamespace()}
+    stream_to_instance_map = {stream_name: AirbyteStream(
+            name=stream_name, namespace=namespace, json_schema={}, supported_sync_modes=[SyncMode.full_refresh, SyncMode.incremental]
+        )
+    }
     state_messages = [AirbyteStateMessage.parse_obj(state_obj) for state_obj in list(input_state)]
     state_manager = ConnectorStateManager(stream_to_instance_map, state_messages)
 

unit_tests/sources/test_integration_source.py

@@ -2,9 +2,7 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-import json
 import os
-from typing import Any, List, Mapping
 from unittest import mock
 from unittest.mock import patch
 
@@ -24,9 +22,9 @@ from unit_tests.sources.fixtures.source_test_fixture import (
     "deployment_mode, url_base, expected_records, expected_error",
     [
         pytest.param("CLOUD", "https://airbyte.com/api/v1/", [], None, id="test_cloud_read_with_public_endpoint"),
-        pytest.param("CLOUD", "http://unsecured.com/api/v1/", [], "system_error", id="test_cloud_read_with_unsecured_url"),
-        pytest.param("CLOUD", "https://172.20.105.99/api/v1/", [], "config_error", id="test_cloud_read_with_private_endpoint"),
-        pytest.param("CLOUD", "https://localhost:80/api/v1/", [], "config_error", id="test_cloud_read_with_localhost"),
+        pytest.param("CLOUD", "http://unsecured.com/api/v1/", [], ValueError, id="test_cloud_read_with_unsecured_url"),
+        pytest.param("CLOUD", "https://172.20.105.99/api/v1/", [], AirbyteTracedException, id="test_cloud_read_with_private_endpoint"),
+        pytest.param("CLOUD", "https://localhost:80/api/v1/", [], AirbyteTracedException, id="test_cloud_read_with_localhost"),
         pytest.param("OSS", "https://airbyte.com/api/v1/", [], None, id="test_oss_read_with_public_endpoint"),
         pytest.param("OSS", "https://172.20.105.99/api/v1/", [], None, id="test_oss_read_with_private_endpoint"),
     ],
@@ -39,10 +37,8 @@ def test_external_request_source(capsys, deployment_mode, url_base, expected_rec
         with mock.patch.object(HttpTestStream, "url_base", url_base):
             args = ["read", "--config", "config.json", "--catalog", "configured_catalog.json"]
             if expected_error:
-                with pytest.raises(AirbyteTracedException):
+                with pytest.raises(expected_error):
                     launch(source, args)
-                    messages = [json.loads(line) for line in capsys.readouterr().out.splitlines()]
-                    assert contains_error_trace_message(messages, expected_error)
             else:
                 launch(source, args)
 
@@ -51,14 +47,14 @@ def test_external_request_source(capsys, deployment_mode, url_base, expected_rec
     "deployment_mode, token_refresh_url, expected_records, expected_error",
     [
         pytest.param("CLOUD", "https://airbyte.com/api/v1/", [], None, id="test_cloud_read_with_public_endpoint"),
-        pytest.param("CLOUD", "http://unsecured.com/api/v1/", [], "system_error", id="test_cloud_read_with_unsecured_url"),
-        pytest.param("CLOUD", "https://172.20.105.99/api/v1/", [], "config_error", id="test_cloud_read_with_private_endpoint"),
+        pytest.param("CLOUD", "http://unsecured.com/api/v1/", [], ValueError, id="test_cloud_read_with_unsecured_url"),
+        pytest.param("CLOUD", "https://172.20.105.99/api/v1/", [], AirbyteTracedException, id="test_cloud_read_with_private_endpoint"),
         pytest.param("OSS", "https://airbyte.com/api/v1/", [], None, id="test_oss_read_with_public_endpoint"),
         pytest.param("OSS", "https://172.20.105.99/api/v1/", [], None, id="test_oss_read_with_private_endpoint"),
     ],
 )
 @patch.object(requests.Session, "send", fixture_mock_send)
-def test_external_oauth_request_source(capsys, deployment_mode, token_refresh_url, expected_records, expected_error):
+def test_external_oauth_request_source(deployment_mode, token_refresh_url, expected_records, expected_error):
     oauth_authenticator = SourceFixtureOauthAuthenticator(
         client_id="nora", client_secret="hae_sung", refresh_token="arthur", token_refresh_endpoint=token_refresh_url
     )
@@ -67,20 +63,7 @@ def test_external_oauth_request_source(capsys, deployment_mode, token_refresh_ur
     with mock.patch.dict(os.environ, {"DEPLOYMENT_MODE": deployment_mode}, clear=False):  # clear=True clears the existing os.environ dict
         args = ["read", "--config", "config.json", "--catalog", "configured_catalog.json"]
         if expected_error:
-            with pytest.raises(AirbyteTracedException):
+            with pytest.raises(expected_error):
                 launch(source, args)
-                messages = [json.loads(line) for line in capsys.readouterr().out.splitlines()]
-                assert contains_error_trace_message(messages, expected_error)
         else:
             launch(source, args)
-
-
-def contains_error_trace_message(messages: List[Mapping[str, Any]], expected_error: str) -> bool:
-    for message in messages:
-        if message.get("type") != "TRACE":
-            continue
-        elif message.get("trace").get("type") != "ERROR":
-            continue
-        elif message.get("trace").get("error").get("failure_type") == expected_error:
-            return True
-    return False

unit_tests/sources/test_source_read.py

@@ -343,7 +343,7 @@ def test_concurrent_source_yields_the_same_messages_as_abstract_source_when_an_e
     source, concurrent_source = _init_sources([stream_slice_to_partition], state, logger)
     config = {}
     catalog = _create_configured_catalog(source._streams)
-    messages_from_abstract_source = _read_from_source(source, logger, config, catalog, state, AirbyteTracedException)
+    messages_from_abstract_source = _read_from_source(source, logger, config, catalog, state, RuntimeError)
     messages_from_concurrent_source = _read_from_source(concurrent_source, logger, config, catalog, state, RuntimeError)
 
     expected_messages = [

unit_tests/test/mock_http/test_mocker.py

@@ -15,6 +15,7 @@ _ANOTHER_RESPONSE_BODY = "another body"
 _A_RESPONSE = HttpResponse("any response")
 _SOME_QUERY_PARAMS = {"q1": "query value"}
 _SOME_HEADERS = {"h1": "header value"}
+_OTHER_HEADERS = {"h2": "another header value"}
 _SOME_REQUEST_BODY_MAPPING = {"first_field": "first_value", "second_field": 2}
 _SOME_REQUEST_BODY_STR = "some_request_body"
 
@@ -24,13 +25,14 @@ class HttpMockerTest(TestCase):
     def test_given_get_request_match_when_decorate_then_return_response(self, http_mocker):
         http_mocker.get(
             HttpRequest(_A_URL, _SOME_QUERY_PARAMS, _SOME_HEADERS),
-            HttpResponse(_A_RESPONSE_BODY, 474),
+            HttpResponse(_A_RESPONSE_BODY, 474, _OTHER_HEADERS),
         )
 
         response = requests.get(_A_URL, params=_SOME_QUERY_PARAMS, headers=_SOME_HEADERS)
 
         assert response.text == _A_RESPONSE_BODY
         assert response.status_code == 474
+        assert response.headers == _OTHER_HEADERS
 
     @HttpMocker()
     def test_given_loose_headers_matching_when_decorate_then_match(self, http_mocker):