Merge pull request #69 from sigmavirus24/feature/61

Add a use_cassette decorator

betamax/decorator.py

@@ -0,0 +1,59 @@
+import functools
+import unittest
+
+import requests
+
+from . import recorder
+
+
+def use_cassette(cassette_name, cassette_library_dir=None,
+                 default_cassette_options={}, **use_cassette_kwargs):
+    """Provide a Betamax-wrapped Session for convenience.
+
+    .. versionadded:: 0.5.0
+
+    This decorator can be used to get a plain Session that has been wrapped in
+    Betamax. For example,
+
+    .. code-block:: python
+
+        from betamax.decorator import use_cassette
+
+        @use_cassette('example-decorator', cassette_library_dir='.')
+        def test_get(session):
+            # do things with session
+
+    :param str cassette_name:
+        Name of the cassette file in which interactions will be stored.
+    :param str cassette_library_dir:
+        Directory in which cassette files will be stored.
+    :param dict default_cassette_options:
+        Dictionary of default cassette options to set for the cassette used
+        when recording these interactions.
+    :param \*\*use_cassette_kwargs:
+        Keyword arguments passed to :meth:`~betamax.Betamax.use_cassette`
+    """
+    def actual_decorator(func):
+        @functools.wraps(func)
+        def test_wrapper(*args, **kwargs):
+            session = requests.Session()
+            recr = recorder.Betamax(
+                session=session,
+                cassette_library_dir=cassette_library_dir,
+                default_cassette_options=default_cassette_options
+            )
+
+            if args:
+                fst, args = args[0], args[1:]
+                if isinstance(fst, unittest.TestCase):
+                    args = (fst, session) + args
+                else:
+                    args = (session, fst) + args
+            else:
+                args = (session,)
+
+            with recr.use_cassette(cassette_name, **use_cassette_kwargs):
+                func(*args, **kwargs)
+
+        return test_wrapper
+    return actual_decorator

docs/api.rst

@@ -6,6 +6,8 @@ API
 .. autoclass:: Betamax
     :members:
 
+.. autofunction:: betamax.decorator.use_cassette
+
 .. autoclass:: betamax.configure.Configuration
     :members:
 
@@ -82,6 +84,26 @@ On the other hand, this will raise an exception:
         r = s.post("https://httpbin.org/post",
                    data={"key": "value"})
 
+
+Finally, we can also use a decorator in order to simplify things:
+
+.. code-block:: python
+
+    import unittest
+
+    from betamax.decorator import use_cassette
+
+    class TestExample(unittest.TestCase):
+        @use_cassette('example', cassette_library_dir='cassettes')
+        def test_example(self, session):
+            session.get('https://httpbin.org/get')
+
+
+    # Or if you're using something like py.test
+    @use_cassette('example', cassette_library_dir='cassettes')
+    def test_example_pytest(session):
+        session.get('https://httpbin.org/get')
+
 .. _opinions:
 
 Opinions at Work

tests/unit/test_decorator.py

@@ -0,0 +1,39 @@
+try:
+    from unittest import mock
+except ImportError:
+    import mock
+
+import betamax
+from betamax.decorator import use_cassette
+
+
+@mock.patch('betamax.recorder.Betamax', autospec=True)
+def test_wraps_session(Betamax):
+    # This needs to be a magic mock so it will mock __exit__
+    recorder = mock.MagicMock(spec=betamax.Betamax)
+    recorder.use_cassette.return_value = recorder
+    Betamax.return_value = recorder
+
+    @use_cassette('foo', cassette_library_dir='fizbarbogus')
+    def _test(session):
+        pass
+
+    _test()
+    Betamax.assert_called_once_with(
+        session=mock.ANY,
+        cassette_library_dir='fizbarbogus',
+        default_cassette_options={}
+    )
+    recorder.use_cassette.assert_called_once_with('foo')
+
+
+@mock.patch('betamax.recorder.Betamax', autospec=True)
+@mock.patch('requests.Session')
+def test_creates_a_new_session(Session, Betamax):
+    @use_cassette('foo', cassette_library_dir='dir')
+    def _test(session):
+        pass
+
+    _test()
+
+    assert Session.call_count == 1