diff --git a/airbyte_cdk/test/mock_http/mocker.py b/airbyte_cdk/test/mock_http/mocker.py index cd1b1f9a7..276f0db09 100644 --- a/airbyte_cdk/test/mock_http/mocker.py +++ b/airbyte_cdk/test/mock_http/mocker.py @@ -17,6 +17,7 @@ class SupportedHttpMethods(str, Enum): GET = "get" PATCH = "patch" POST = "post" + PUT = "put" DELETE = "delete" @@ -77,7 +78,7 @@ def _mock_request_method( additional_matcher=self._matches_wrapper(matcher), response_list=[ { - "text": response.body, + self._get_body_field(response): response.body, "status_code": response.status_code, "headers": response.headers, } @@ -85,6 +86,10 @@ def _mock_request_method( ], ) + @staticmethod + def _get_body_field(response: HttpResponse) -> str: + return "text" if isinstance(response.body, str) else "content" + def get(self, request: HttpRequest, responses: Union[HttpResponse, List[HttpResponse]]) -> None: self._mock_request_method(SupportedHttpMethods.GET, request, responses) @@ -98,6 +103,9 @@ def post( ) -> None: self._mock_request_method(SupportedHttpMethods.POST, request, responses) + def put(self, request: HttpRequest, responses: Union[HttpResponse, List[HttpResponse]]) -> None: + self._mock_request_method(SupportedHttpMethods.PUT, request, responses) + def delete( self, request: HttpRequest, responses: Union[HttpResponse, List[HttpResponse]] ) -> None: diff --git a/airbyte_cdk/test/mock_http/response.py b/airbyte_cdk/test/mock_http/response.py index 848be55a0..fefe762e9 100644 --- a/airbyte_cdk/test/mock_http/response.py +++ b/airbyte_cdk/test/mock_http/response.py @@ -1,19 +1,22 @@ # Copyright (c) 2023 Airbyte, Inc., all rights reserved. from types import MappingProxyType -from typing import Mapping +from typing import Mapping, Union class HttpResponse: def __init__( - self, body: str, status_code: int = 200, headers: Mapping[str, str] = MappingProxyType({}) + self, + body: Union[str, bytes], + status_code: int = 200, + headers: Mapping[str, str] = MappingProxyType({}), ): self._body = body self._status_code = status_code self._headers = headers @property - def body(self) -> str: + def body(self) -> Union[str, bytes]: return self._body @property diff --git a/unit_tests/test/mock_http/test_mocker.py b/unit_tests/test/mock_http/test_mocker.py index 4ff5afe01..6a50a30c2 100644 --- a/unit_tests/test/mock_http/test_mocker.py +++ b/unit_tests/test/mock_http/test_mocker.py @@ -1,5 +1,5 @@ # Copyright (c) 2023 Airbyte, Inc., all rights reserved. - +import gzip from unittest import TestCase import pytest @@ -75,6 +75,34 @@ def test_given_post_request_match_when_decorate_then_return_response(self, http_ assert response.text == _A_RESPONSE_BODY assert response.status_code == 474 + @HttpMocker() + def test_given_put_request_match_when_decorate_then_return_response(self, http_mocker): + http_mocker.put( + HttpRequest(_A_URL, _SOME_QUERY_PARAMS, _SOME_HEADERS, _SOME_REQUEST_BODY_STR), + HttpResponse(_A_RESPONSE_BODY, 474), + ) + + response = requests.put( + _A_URL, params=_SOME_QUERY_PARAMS, headers=_SOME_HEADERS, data=_SOME_REQUEST_BODY_STR + ) + + assert response.text == _A_RESPONSE_BODY + assert response.status_code == 474 + + @HttpMocker() + def test_given_response_in_bytes_when_decorate_then_match(self, http_mocker): + response_content = gzip.compress(bytes("This is the response within the gzip", "utf-8")) + http_mocker.put( + HttpRequest(_A_URL, _SOME_QUERY_PARAMS, _SOME_HEADERS, _SOME_REQUEST_BODY_STR), + HttpResponse(response_content, 474), + ) + + response = requests.put( + _A_URL, params=_SOME_QUERY_PARAMS, headers=_SOME_HEADERS, data=_SOME_REQUEST_BODY_STR + ) + + assert response.content == response_content + @HttpMocker() def test_given_multiple_responses_when_decorate_get_request_then_return_response( self, http_mocker