Repository: betamaxpy/betamax

Issue:
Interest in additional hooks?
A recently merged change in [PRAW](/praw-dev/praw) (which uses Betamax for its integration tests) hacks in an "after init" hook:

```python
def add_init_hook(original_init):
    """Wrap an __init__ method to also call some hooks."""

    @wraps(original_init)
    def wrapper(self, *args, **kwargs):
        original_init(self, *args, **kwargs)
        dispatch_hooks("after_init", self)

    return wrapper


Cassette.__init__ = add_init_hook(Cassette.__init__)


def init_hook(cassette):
    if cassette.is_recording():
        pytest.set_up_record()  # dynamically defined in __init__.py


Cassette.hooks["after_init"].append(init_hook)
```

[source code in PRAW repo](https://github.com/praw-dev/praw/blob/edf9650a415bf484b3d6d886dca65517db141ac5/tests/conftest.py#L93-L104)

For PRAW's tests, this hook is useful because it triggers at the time `.use_cassette()` is called, which in the case of PRAW's tests signifies that network requests will happen soon but they haven't started yet. This hook allows us to modify the test environment slightly before any requests are made, based on whether a certain cassette has been recorded or not.

Rather than having to monkeypatch Betamax, it would be better for PRAW if this hook existed natively in Betamax. **Would this be a welcome feature in Betamax?** If so, I'll be happy to write the feature.

Comments:

Patch:
diff --git a/AUTHORS.rst b/AUTHORS.rst
index 8bbdd19..078debf 100644
--- a/AUTHORS.rst
+++ b/AUTHORS.rst
@@ -19,3 +19,4 @@ Contributors
 - Marc Abramowitz (@msabramo)
 - Bryce Boe <bbzbryce@gmail.com> (@bboe)
 - Alex Richard-Hoyling <@arhoyling)
+- Joey RH <jarhill0@gmail.com> (@jarhill0)
diff --git a/docs/source/configuring.rst b/docs/source/configuring.rst
index 93e70e4..361bec6 100644
--- a/docs/source/configuring.rst
+++ b/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
 ``````````````````````````
diff --git a/src/betamax/configure.py b/src/betamax/configure.py
index 1eca77c..b065eff 100644
--- a/src/betamax/configure.py
+++ b/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 @@ class Configuration(object):
         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 @@ class Configuration(object):
         """
         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."""
diff --git a/src/betamax/recorder.py b/src/betamax/recorder.py
index 2a4b069..d89b8bd 100644
--- a/src/betamax/recorder.py
+++ b/src/betamax/recorder.py
@@ -123,10 +123,13 @@ class Betamax(object):
         """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 @@ class Betamax(object):
             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)


Test patch:
diff --git a/tests/integration/test_hooks.py b/tests/integration/test_hooks.py
index 8378e5d..4cfa7f8 100644
--- a/tests/integration/test_hooks.py
+++ b/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:
diff --git a/tests/unit/test_configure.py b/tests/unit/test_configure.py
index c39de63..4f00b93 100644
--- a/tests/unit/test_configure.py
+++ b/tests/unit/test_configure.py
@@ -4,6 +4,7 @@ import unittest
 
 from betamax.configure import Configuration
 from betamax.cassette import Cassette
+from betamax.recorder import Betamax
 
 
 class TestConfiguration(unittest.TestCase):
@@ -14,6 +15,7 @@ class TestConfiguration(unittest.TestCase):
         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 @@ class TestConfiguration(unittest.TestCase):
         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 @@ class TestConfiguration(unittest.TestCase):
         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])