Add unit test support for journeys

changes:
- introduce a new mite.utils.sleep function (so that tests donʼt sleep
  for a long time)
- change the base class of HandledException to work around a bug in
  pytest-asyncio (pytest-dev/pytest-asyncio#126)
- better async exception handling
- fix event loop handling in mite_http
- new mite.test module

mite/__init__.py

@@ -11,9 +11,10 @@
 from .exceptions import MiteError
 from .context import Context
 from .runner import RunnerConfig
+import mite.utils
 
-from .exceptions import MiteError
 
+# TODO: move to test.py?
 def test_context(extensions=('http',), **config):
     runner_config = RunnerConfig()
     runner_config._update(config.items())
@@ -37,7 +38,7 @@ async def __aexit__(self, *args):
             self._loop = asyncio.get_event_loop()
         sleep_time = self._sleep_time()
         if sleep_time > 0:
-            await asyncio.sleep(sleep_time, loop=self._loop)
+            await mite.utils.sleep(sleep_time, loop=self._loop)
 
 
 def ensure_fixed_separation(separation, loop=None):

mite/context.py

@@ -13,7 +13,7 @@
 logger = logging.getLogger(__name__)
 
 
-class HandledException(BaseException):
+class HandledException(Exception):
     def __init__(self, original_exception, original_tb):
         self.original_exception = original_exception
         self.original_tb = original_tb
@@ -77,7 +77,7 @@ async def transaction(self, name):
                 drop_to_debugger(traceback)
                 return
             else:
-                raise new_exc
+                raise new_exc from exception_val
         finally:
             await self._end_transaction()
 

mite/test.py

@@ -0,0 +1,101 @@
+from .context import Context
+from .config import default_config_loader
+from collections import defaultdict
+import pytest
+import asyncio
+import mite.utils as utils
+
+
+class _InterceptHttp:
+    def __init__(self, http, obj):
+        http._parent = self
+        self._obj = obj
+        object.__setattr__(self, "http", http)
+        self._old_http = None
+
+    def __setattr__(self, name, val):
+        if name == "http":
+            object.__setattr__(self, "_old_http", val)
+        elif name in ("_obj", "_old_http"):
+            object.__setattr__(self, name, val)
+        else:
+            self._obj.__setattr__(name, val)
+
+    def __getattribute__(self, name):
+        try:
+            return object.__getattribute__(self, name)
+        except AttributeError:
+            return object.__getattribute__(self, "_obj").__getattribute__(name)
+
+    def __delattr__(self, name):
+        if name == "http":
+            object.__delattr__(self, "_old_http")
+            return
+        self._obj.__delattr__(name)
+
+
+class _NewHttp:
+    def __init__(self, requests):
+        self._requests = requests
+
+    async def get(self, *args, **kwargs):
+        r = await self._parent._old_http.get(*args, **kwargs)
+        self._requests['get'].append(r)
+        return r
+
+    async def post(self, *args, **kwargs):
+        r = await self._parent._old_http.post(*args, **kwargs)
+        self._requests['post'].append(r)
+        return r
+
+    async def delete(self, *args, **kwargs):
+        r = await self._parent._old_http.delete(*args, **kwargs)
+        self._requests['delete'].append(r)
+        return r
+
+    async def put(self, *args, **kwargs):
+        r = await self._parent._old_http.put(*args, **kwargs)
+        self._requests['put'].append(r)
+        return r
+
+
+def http_spy(journey):
+    async def wrapper(ctx, *args):
+        requests = defaultdict(list)
+        # The reason for doing this in terms of this (admittedly complicated)
+        # dance is the following: we want to be able to say
+        # "http_spy(journey)".  The journey function is already wrapped with a
+        # decorator (namely mite_http) which, when it is called with a ctx
+        # argument, will inject a http attribute on that ctx.  We want to pass
+        # our own http argument instead.  If we just add it now, it will be
+        # overwritten by mite_http.  We can't "undecorate" the journey, at
+        # least not without intrusive modifications to the mite codebase.  So
+        # we use the above nasty hack with __getattribute__ and friends.  An
+        # alternative might be to investigate a prope mocking/spying library
+        # in python...
+        spy = _InterceptHttp(_NewHttp(requests), ctx)
+        await journey(spy, *args)
+        return {'http': requests}
+    return wrapper
+
+
+async def run_single_journey(journey, datapool=None):
+    messages = []
+
+    async def _send(message, **kwargs):
+        messages.append((message, kwargs))
+
+    config = default_config_loader()
+
+    ctx = Context(_send, config)
+
+    if datapool is not None:
+        dpi = datapool.checkout()
+        result = await journey(ctx, *dpi.data)
+    else:
+        result = await journey(ctx)
+
+    if result is None:
+        result = {}
+    result['messages'] = messages
+    return result

mite/utils.py

@@ -1,6 +1,6 @@
 import msgpack
 import importlib
-
+import asyncio
 
 _msg_unpacker = msgpack.Unpacker(raw=False, use_list=False)
 
@@ -17,3 +17,7 @@ def unpack_msg(msg):
 def spec_import(spec):
     module, attr = spec.split(':', 1)
     return getattr(importlib.import_module(module), attr)
+
+
+async def sleep(delay, always=False, **kwargs):
+    await asyncio.sleep(delay, **kwargs)

mite_http/__init__.py

@@ -1,6 +1,7 @@
 from collections import deque
 from acurl import EventLoop
 import logging
+import asyncio
 
 from mite.compat import asynccontextmanager
 
@@ -54,9 +55,15 @@ async def _checkin(self, session):
 
 
 def get_session_pool():
-    if not hasattr(get_session_pool, '_session_pool'):
-        get_session_pool._session_pool = SessionPool()
-    return get_session_pool._session_pool
+    # We memoize the function by event loop.  This is because, in unit tests,
+    # there are multiple event loops in circulation.
+    try:
+        return get_session_pool._session_pools[asyncio.get_event_loop()]
+    except KeyError:
+        sp = SessionPool()
+        get_session_pool._session_pools[asyncio.get_event_loop()] = sp
+        return sp
+get_session_pool._session_pools = {}
 
 
 def mite_http(func):