Merge pull request #179 from jarhill0/pre-post-hooks

Add cassette_load and cassette_eject hooks

AUTHORS.rst

@@ -19,3 +19,4 @@ Contributors
 - Marc Abramowitz (@msabramo)
 - Bryce Boe <[email protected]> (@bboe)
 - Alex Richard-Hoyling <@arhoyling)
+- Joey RH <[email protected]> (@jarhill0)

docs/source/configuring.rst

@@ -241,6 +241,27 @@ have our hook, we need merely register it like so:
 
 And we no longer need to worry about leaking sensitive data.
 
+In addition to the ``before_record`` and ``before_playback`` hooks,
+version 0.9.0 of Betamax adds :meth:`.after_start` and :meth:`.before_stop`
+hooks. These two hooks both will pass the current
+:class:`~betamax.cassette.cassette.Cassette` to the callback function provided.
+Register these hooks like so:
+
+.. code-block:: python
+
+    def hook(cassette):
+        if cassette.is_recording():
+            print("This cassette is recording!")
+
+    # Either
+    config.after_start(callback=hook)
+    # Or
+    config.before_stop(callback=hook)
+
+These hooks are useful for performing configuration actions external to Betamax
+at the time Betamax is invoked, such as setting up correct authentication to
+an API so that the recording will not encounter any errors.
+
 
 Setting default serializer
 ``````````````````````````

src/betamax/configure.py

@@ -1,8 +1,9 @@
+from collections import defaultdict
+
 from .cassette import Cassette
 
 
 class Configuration(object):
-
     """This object acts as a proxy to configure different parts of Betamax.
 
     You should only ever encounter this object when configuring the library as
@@ -20,6 +21,7 @@ class Configuration(object):
     """
 
     CASSETTE_LIBRARY_DIR = 'vcr/cassettes'
+    recording_hooks = defaultdict(list)
 
     def __enter__(self):
         return self
@@ -33,6 +35,26 @@ def __setattr__(self, prop, value):
         else:
             super(Configuration, self).__setattr__(prop, value)
 
+    def after_start(self, callback=None):
+        """Register a function to call after Betamax is started.
+
+        Example usage:
+
+        .. code-block:: python
+
+            def on_betamax_start(cassette):
+                if cassette.is_recording():
+                    print("Setting up authentication...")
+
+            with Betamax.configure() as config:
+                config.cassette_load(callback=on_cassette_load)
+
+        :param callable callback:
+            The function which accepts a cassette and might mutate
+            it before returning.
+        """
+        self.recording_hooks['after_start'].append(callback)
+
     def before_playback(self, tag=None, callback=None):
         """Register a function to call before playing back an interaction.
 
@@ -79,6 +101,26 @@ def before_record(interaction, cassette):
         """
         Cassette.hooks['before_record'].append(callback)
 
+    def before_stop(self, callback=None):
+        """Register a function to call before Betamax stops.
+
+        Example usage:
+
+        .. code-block:: python
+
+            def on_betamax_stop(cassette):
+                if not cassette.is_recording():
+                    print("Playback completed.")
+
+            with Betamax.configure() as config:
+                config.cassette_eject(callback=on_betamax_stop)
+
+        :param callable callback:
+            The function which accepts a cassette and might mutate
+            it before returning.
+        """
+        self.recording_hooks['before_stop'].append(callback)
+
     @property
     def cassette_library_dir(self):
         """Retrieve and set the directory to store the cassettes in."""

src/betamax/recorder.py

@@ -123,10 +123,13 @@ def start(self):
         """Start recording or replaying interactions."""
         for k in self.http_adapters:
             self.session.mount(k, self.betamax_adapter)
+        dispatch_hooks('after_start', self.betamax_adapter.cassette)
 
     # ■
     def stop(self):
         """Stop recording or replaying interactions."""
+        dispatch_hooks('before_stop', self.betamax_adapter.cassette)
+
         # No need to keep the cassette in memory any longer.
         self.betamax_adapter.eject_cassette()
         # On exit, we no longer wish to use our adapter and we want the
@@ -166,3 +169,10 @@ def use_cassette(self, cassette_name, **kwargs):
             raise ValueError('Cassette must have a valid name and may not be'
                              ' None.')
         return self
+
+
+def dispatch_hooks(hook_name, *args):
+    """Dispatch registered hooks."""
+    hooks = Configuration.recording_hooks[hook_name]
+    for hook in hooks:
+        hook(*args)

tests/integration/test_hooks.py

@@ -17,12 +17,60 @@ def preplayback_hook(interaction, cassette):
     interaction.data['response']['headers']['Betamax-Fake-Header'] = 'temp'
 
 
+class Counter(object):
+    def __init__(self):
+        self.value = 0
+
+    def increment(self, cassette):
+        self.value += 1
+
+
 class TestHooks(helper.IntegrationHelper):
     def tearDown(self):
         super(TestHooks, self).tearDown()
         # Clear out the hooks
+        betamax.configure.Configuration.recording_hooks.pop('after_start', None)
         betamax.cassette.Cassette.hooks.pop('before_record', None)
         betamax.cassette.Cassette.hooks.pop('before_playback', None)
+        betamax.configure.Configuration.recording_hooks.pop('before_stop', None)
+
+    def test_post_start_hook(self):
+        start_count = Counter()
+        with betamax.Betamax.configure() as config:
+            config.after_start(callback=start_count.increment)
+
+        recorder = betamax.Betamax(self.session)
+
+        assert start_count.value == 0
+        with recorder.use_cassette('after_start_hook'):
+            assert start_count.value == 1
+            self.cassette_path = recorder.current_cassette.cassette_path
+            self.session.get('https://httpbin.org/get')
+
+        assert start_count.value == 1
+        with recorder.use_cassette('after_start_hook', record='none'):
+            assert start_count.value == 2
+            self.session.get('https://httpbin.org/get')
+        assert start_count.value == 2
+
+    def test_pre_stop_hook(self):
+        stop_count = Counter()
+        with betamax.Betamax.configure() as config:
+            config.before_stop(callback=stop_count.increment)
+
+        recorder = betamax.Betamax(self.session)
+
+        assert stop_count.value == 0
+        with recorder.use_cassette('before_stop_hook'):
+            self.cassette_path = recorder.current_cassette.cassette_path
+            self.session.get('https://httpbin.org/get')
+            assert stop_count.value == 0
+        assert stop_count.value == 1
+
+        with recorder.use_cassette('before_stop_hook', record='none'):
+            self.session.get('https://httpbin.org/get')
+            assert stop_count.value == 1
+        assert stop_count.value == 2
 
     def test_prerecord_hook(self):
         with betamax.Betamax.configure() as config:

tests/unit/test_configure.py

@@ -4,6 +4,7 @@
 
 from betamax.configure import Configuration
 from betamax.cassette import Cassette
+from betamax.recorder import Betamax
 
 
 class TestConfiguration(unittest.TestCase):
@@ -14,6 +15,7 @@ def setUp(self):
         self.cassette_dir = Configuration.CASSETTE_LIBRARY_DIR
 
     def tearDown(self):
+        Configuration.recording_hooks = collections.defaultdict(list)
         Cassette.default_cassette_options = self.cassette_options
         Cassette.hooks = collections.defaultdict(list)
         Configuration.CASSETTE_LIBRARY_DIR = self.cassette_dir
@@ -43,6 +45,14 @@ def test_allows_registration_of_placeholders(self):
         assert placeholders[0]['placeholder'] == '<TEST>'
         assert placeholders[0]['replace'] == 'test'
 
+    def test_registers_post_start_hooks(self):
+        c = Configuration()
+        assert Configuration.recording_hooks['after_start'] == []
+        c.after_start(callback=lambda: None)
+        assert Configuration.recording_hooks['after_start'] != []
+        assert len(Configuration.recording_hooks['after_start']) == 1
+        assert callable(Configuration.recording_hooks['after_start'][0])
+
     def test_registers_pre_record_hooks(self):
         c = Configuration()
         assert Cassette.hooks['before_record'] == []
@@ -58,3 +68,11 @@ def test_registers_pre_playback_hooks(self):
         assert Cassette.hooks['before_playback'] != []
         assert len(Cassette.hooks['before_playback']) == 1
         assert callable(Cassette.hooks['before_playback'][0])
+
+    def test_registers_pre_stop_hooks(self):
+        c = Configuration()
+        assert Configuration.recording_hooks['before_stop'] == []
+        c.before_stop(callback=lambda: None)
+        assert Configuration.recording_hooks['before_stop'] != []
+        assert len(Configuration.recording_hooks['before_stop']) == 1
+        assert callable(Configuration.recording_hooks['before_stop'][0])