From c0a22df7ede3261eeb178e91a061994a173ee981 Mon Sep 17 00:00:00 2001 From: Ivan Malison Date: Thu, 8 Jan 2015 10:53:44 -0800 Subject: [PATCH] Add ability to add custom patches to vcr and cassettes. --- tests/unit/test_cassettes.py | 22 +++++++++++++++++++++- tests/unit/test_vcr.py | 16 ++++++++++++++++ vcr/cassette.py | 3 ++- vcr/config.py | 9 +++++++-- vcr/patch.py | 20 +++++++++++++++++--- 5 files changed, 63 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_cassettes.py b/tests/unit/test_cassettes.py index 2cab15f..d68ecea 100644 --- a/tests/unit/test_cassettes.py +++ b/tests/unit/test_cassettes.py @@ -7,8 +7,10 @@ import pytest import yaml from vcr.cassette import Cassette -from vcr.patch import force_reset from vcr.errors import UnhandledHTTPRequestError +from vcr.patch import force_reset +from vcr.stubs import VCRHTTPSConnection + def test_cassette_load(tmpdir): @@ -181,3 +183,21 @@ def test_nesting_context_managers_by_checking_references_of_http_connection(): assert httplib.HTTPConnection is original assert httplib.HTTPConnection is second_cassette_HTTPConnection assert httplib.HTTPConnection is first_cassette_HTTPConnection + + +def test_custom_patchers(): + class Test(object): + attribute = None + with Cassette.use('custom_patches', custom_patches=((Test, 'attribute', VCRHTTPSConnection),)): + assert issubclass(Test.attribute, VCRHTTPSConnection) + assert VCRHTTPSConnection is not Test.attribute + old_attribute = Test.attribute + + with Cassette.use('custom_patches', custom_patches=((Test, 'attribute', VCRHTTPSConnection),)): + assert issubclass(Test.attribute, VCRHTTPSConnection) + assert VCRHTTPSConnection is not Test.attribute + assert Test.attribute is not old_attribute + + assert issubclass(Test.attribute, VCRHTTPSConnection) + assert VCRHTTPSConnection is not Test.attribute + assert Test.attribute is old_attribute diff --git a/tests/unit/test_vcr.py b/tests/unit/test_vcr.py index 33e17a0..0c1c6c7 100644 --- a/tests/unit/test_vcr.py +++ b/tests/unit/test_vcr.py @@ -3,6 +3,7 @@ import pytest from vcr import VCR, use_cassette from vcr.request import Request +from vcr.stubs import VCRHTTPSConnection def test_vcr_use_cassette(): @@ -74,3 +75,18 @@ def test_fixtures_with_use_cassette(random_fixture): # fixtures. It is admittedly a bit strange because the test would never even # run if the relevant feature were broken. pass + + +def test_custom_patchers(): + class Test(object): + attribute = None + attribute2 = None + test_vcr = VCR(custom_patches=((Test, 'attribute', VCRHTTPSConnection),)) + with test_vcr.use_cassette('custom_patches'): + assert issubclass(Test.attribute, VCRHTTPSConnection) + assert VCRHTTPSConnection is not Test.attribute + + with test_vcr.use_cassette('custom_patches', custom_patches=((Test, 'attribute2', VCRHTTPSConnection),)): + assert issubclass(Test.attribute, VCRHTTPSConnection) + assert VCRHTTPSConnection is not Test.attribute + assert Test.attribute is Test.attribute2 diff --git a/vcr/cassette.py b/vcr/cassette.py index 7a8a4dd..d8c63c9 100644 --- a/vcr/cassette.py +++ b/vcr/cassette.py @@ -87,7 +87,7 @@ class Cassette(object): match_on=(uri, method), filter_headers=(), filter_query_parameters=(), before_record_request=None, before_record_response=None, ignore_hosts=(), - ignore_localhost=()): + ignore_localhost=(), custom_patches=()): self._path = path self._serializer = serializer self._match_on = match_on @@ -100,6 +100,7 @@ class Cassette(object): self.dirty = False self.rewound = False self.record_mode = record_mode + self.custom_patches = custom_patches @property def play_count(self): diff --git a/vcr/config.py b/vcr/config.py index e8ff7a8..88ae535 100644 --- a/vcr/config.py +++ b/vcr/config.py @@ -12,7 +12,7 @@ from . import filters class VCR(object): def __init__(self, serializer='yaml', cassette_library_dir=None, - record_mode="once", filter_headers=(), + record_mode="once", filter_headers=(), custom_patches=(), filter_query_parameters=(), before_record_request=None, before_record_response=None, ignore_hosts=(), match_on=('method', 'scheme', 'host', 'port', 'path', 'query',), @@ -43,6 +43,7 @@ class VCR(object): self.before_record_response = before_record_response self.ignore_hosts = ignore_hosts self.ignore_localhost = ignore_localhost + self._custom_patches = tuple(custom_patches) def _get_serializer(self, serializer_name): try: @@ -68,6 +69,9 @@ class VCR(object): def use_cassette(self, path, with_current_defaults=False, **kwargs): if with_current_defaults: return Cassette.use(path, self.get_path_and_merged_config(path, **kwargs)) + # This is made a function that evaluates every time a cassette is made so that + # changes that are made to this VCR instance that occur AFTER the use_cassette + # decorator is applied still affect subsequent calls to the decorated function. args_getter = functools.partial(self.get_path_and_merged_config, path, **kwargs) return Cassette.use_arg_getter(args_getter) @@ -86,7 +90,8 @@ class VCR(object): 'match_on': self._get_matchers(matcher_names), 'record_mode': kwargs.get('record_mode', self.record_mode), 'before_record_request': self._build_before_record_request(kwargs), - 'before_record_response': self._build_before_record_response(kwargs) + 'before_record_response': self._build_before_record_response(kwargs), + 'custom_patches': self._custom_patches + kwargs.get('custom_patches', ()) } return path, merged_config diff --git a/vcr/patch.py b/vcr/patch.py index 7e538a9..9408397 100644 --- a/vcr/patch.py +++ b/vcr/patch.py @@ -69,9 +69,12 @@ class CassettePatcherBuilder(object): self._class_to_cassette_subclass = {} def build(self): - return itertools.chain(self._httplib(), self._requests(), - self._urllib3(), self._httplib2(), - self._boto()) + return itertools.chain( + self._httplib(), self._requests(), self._urllib3(), self._httplib2(), + self._boto(), self._build_patchers_from_mock_triples( + self._cassette.custom_patches + ) + ) def _build_patchers_from_mock_triples(self, mock_triples): for args in mock_triples: @@ -88,6 +91,17 @@ class CassettePatcherBuilder(object): replacement_class)) def _recursively_apply_get_cassette_subclass(self, replacement_dict_or_obj): + """One of the subtleties of this class is that it does not directly + replace HTTPSConnection with VCRRequestsHTTPSConnection, but a + subclass of this class that has cassette assigned to the + appropriate value. This behavior is necessary to properly + support nested cassette contexts + + This function exists to ensure that we use the same class + object (reference) to patch everything that replaces + VCRRequestHTTP[S]Connection, but that we can talk about + patching them with the raw references instead. + """ if isinstance(replacement_dict_or_obj, dict): for key, replacement_obj in replacement_dict_or_obj.items(): replacement_obj = self._recursively_apply_get_cassette_subclass(