Add gevent.util.assert_switches

Based on refactoring existing code in _monitor.py

Flexibly allows checking for any switching, switches that take longer
than a set amount of time, or switches that exceed a time limit
between getting back to the hub.

Fixes #1182.

CHANGES.rst

@@ -26,6 +26,9 @@
 
 - The long-deprecated and undocumented module ``gevent.wsgi`` was removed.
 
+- Add `gevent.util.assert_switches` to build on the monitoring
+  functions. Fixes :issue:`1182`.
+
 1.3b1 (2018-04-13)
 ==================
 

doc/monitoring.rst

@@ -27,6 +27,9 @@ interval. When such a blocking greenlet is detected, it will print
 :attr:`~gevent.hub.Hub.exception_stream`. It will also emit the
 :class:`gevent.events.EventLoopBlocked` event.
 
+.. seealso:: :func:`gevent.util.assert_switches`
+
+   For a scoped version of this.
 
 Memory Usage
 ------------

src/gevent/util.py

@@ -23,6 +23,7 @@ __all__ = [
     'print_run_info',
     'GreenletTree',
     'wrap_errors',
+    'assert_switches',
 ]
 
 # PyPy is very slow at formatting stacks
@@ -507,3 +508,85 @@ class GreenletTree(object):
         Returns the `GreenletTree` for the current thread.
         """
         return cls._forest()[1]
+
+class _FailedToSwitch(AssertionError):
+    pass
+
+class assert_switches(object):
+    """
+    A context manager for ensuring a block of code switches greenlets.
+
+    This performs a similar function as the :doc:`monitoring thread
+    </monitoring>`, but the scope is limited to the body of the with
+    statement. If the code within the body doesn't yield to the hub
+    (and doesn't raise an exception), then upon exiting the
+    context manager an :exc:`AssertionError` will be raised.
+
+    This is useful in unit tests and for debugging purposes.
+
+    :keyword float max_blocking_time: If given, the body is allowed
+        to block for up to this many fractional seconds before
+        an error is raised.
+    :keyword bool hub_only: If True, then *max_blocking_time* only
+        refers to the amount of time spent between switches into the
+        hub. If False, then it refers to the maximum time between
+        *any* switches. If *max_blocking_time* is not given, has no
+        effect.
+
+    Example::
+
+        # This will always raise an exception: nothing switched
+        with assert_switches():
+            pass
+
+        # This will never raise an exception; nothing switched,
+        # but it happened very fast
+        with assert_switches(max_blocking_time=1.0):
+            pass
+
+    .. versionadded:: 1.3
+    """
+
+    hub = None
+    tracer = None
+
+
+    def __init__(self, max_blocking_time=None, hub_only=False):
+        self.max_blocking_time = max_blocking_time
+        self.hub_only = hub_only
+
+    def __enter__(self):
+        from gevent import get_hub
+        from gevent import _monitor
+
+        self.hub = hub = get_hub()
+
+        # TODO: We could optimize this to use the GreenletTracer
+        # installed by the monitoring thread, if there is one.
+        # As it is, we will chain trace calls back to it.
+        if not self.max_blocking_time:
+            self.tracer = _monitor.GreenletTracer()
+        elif self.hub_only:
+            self.tracer = _monitor.HubSwitchTracer(hub, self.max_blocking_time)
+        else:
+            self.tracer = _monitor.MaxSwitchTracer(hub, self.max_blocking_time)
+
+        self.tracer.monitor_current_greenlet_blocking()
+        return self
+
+    def __exit__(self, t, v, tb):
+        self.tracer.kill()
+        hub = self.hub; self.hub = None
+        tracer = self.tracer; self.tracer = None
+
+        # Only check if there was no exception raised, we
+        # don't want to hide anything
+        if t is not None:
+            return
+
+
+        did_block = tracer.did_block_hub(hub)
+        if did_block:
+            active_greenlet = did_block[1]
+            report_lines = tracer.did_block_hub_report(hub, active_greenlet, {})
+            raise _FailedToSwitch('\n'.join(report_lines))

src/greentest/test___monitor.py

@@ -50,7 +50,7 @@ class _AbstractTestPeriodicMonitoringThread(object):
     def tearDown(self):
         monitor.start_new_thread = self._orig_start_new_thread
         monitor.thread_sleep = self._orig_thread_sleep
-        prev = self.pmt.previous_trace_function
+        prev = self.pmt._greenlet_tracer.previous_trace_function
         self.pmt.kill()
         assert gettrace() is prev, (gettrace(), prev)
         settrace(None)
@@ -62,7 +62,7 @@ class TestPeriodicMonitoringThread(_AbstractTestPeriodicMonitoringThread,
 
     def test_constructor(self):
         self.assertEqual(0xDEADBEEF, self.pmt.monitor_thread_ident)
-        self.assertEqual(gettrace(), self.pmt.greenlet_trace)
+        self.assertEqual(gettrace(), self.pmt._greenlet_tracer)
 
     def test_hub_wref(self):
         self.assertIs(self.hub, self.pmt.hub)
@@ -178,31 +178,31 @@ class TestPeriodicMonitorBlocking(_AbstractTestPeriodicMonitoringThread,
         settrace(f)
 
         self.pmt = monitor.PeriodicMonitoringThread(self.hub)
-        self.assertEqual(gettrace(), self.pmt.greenlet_trace)
-        self.assertIs(self.pmt.previous_trace_function, f)
+        self.assertEqual(gettrace(), self.pmt._greenlet_tracer)
+        self.assertIs(self.pmt._greenlet_tracer.previous_trace_function, f)
 
-        self.pmt.greenlet_trace('event', 'args')
+        self.pmt._greenlet_tracer('event', 'args')
 
         self.assertEqual([('event', 'args')], called)
 
-    def test_greenlet_trace(self):
-        self.assertEqual(0, self.pmt._greenlet_switch_counter)
+    def test__greenlet_tracer(self):
+        self.assertEqual(0, self.pmt._greenlet_tracer.greenlet_switch_counter)
         # Unknown event still counts as a switch (should it?)
-        self.pmt.greenlet_trace('unknown', None)
-        self.assertEqual(1, self.pmt._greenlet_switch_counter)
-        self.assertIsNone(self.pmt._active_greenlet)
+        self.pmt._greenlet_tracer('unknown', None)
+        self.assertEqual(1, self.pmt._greenlet_tracer.greenlet_switch_counter)
+        self.assertIsNone(self.pmt._greenlet_tracer.active_greenlet)
 
         origin = object()
         target = object()
 
-        self.pmt.greenlet_trace('switch', (origin, target))
-        self.assertEqual(2, self.pmt._greenlet_switch_counter)
-        self.assertIs(target, self.pmt._active_greenlet)
+        self.pmt._greenlet_tracer('switch', (origin, target))
+        self.assertEqual(2, self.pmt._greenlet_tracer.greenlet_switch_counter)
+        self.assertIs(target, self.pmt._greenlet_tracer.active_greenlet)
 
         # Unknown event removes active greenlet
-        self.pmt.greenlet_trace('unknown', self)
-        self.assertEqual(3, self.pmt._greenlet_switch_counter)
-        self.assertIsNone(self.pmt._active_greenlet)
+        self.pmt._greenlet_tracer('unknown', self)
+        self.assertEqual(3, self.pmt._greenlet_tracer.greenlet_switch_counter)
+        self.assertIsNone(self.pmt._greenlet_tracer.active_greenlet)
 
     def test_monitor_blocking(self):
         # Initially there's no active greenlet and no switches,
@@ -218,7 +218,7 @@ class TestPeriodicMonitorBlocking(_AbstractTestPeriodicMonitoringThread,
         # Give it an active greenlet
         origin = object()
         target = object()
-        self.pmt.greenlet_trace('switch', (origin, target))
+        self.pmt._greenlet_tracer('switch', (origin, target))
 
         # We've switched, so we're not blocked
         self.assertFalse(self.pmt.monitor_blocking(self.hub))

src/greentest/test__util.py

@@ -7,6 +7,7 @@ from __future__ import division
 from __future__ import print_function
 
 import gc
+import unittest
 
 import greentest
 
@@ -208,5 +209,64 @@ class TestTree(greentest.TestCase):
 
         self.assertEqual(expected, value)
 
+
+class TestAssertSwitches(unittest.TestCase):
+
+    def test_time_sleep(self):
+        # A real blocking function
+        from time import sleep
+        with self.assertRaises(util._FailedToSwitch):
+            with util.assert_switches():
+                sleep(0.001)
+
+        # Supply a max allowed and exceed it
+        with self.assertRaises(util._FailedToSwitch):
+            with util.assert_switches(0.001):
+                sleep(0.1)
+
+
+        # Stay within it, but don't switch to the hub
+        with self.assertRaises(util._FailedToSwitch):
+            with util.assert_switches(0.001, hub_only=True):
+                sleep(0)
+
+        # Stay within it, and we only watch for any switch
+        with util.assert_switches(0.001, hub_only=False):
+            sleep(0)
+
+
+    def test_no_switches_no_function(self):
+        # No blocking time given, no switch performed: exception
+        with self.assertRaises(util._FailedToSwitch):
+            with util.assert_switches():
+                pass
+
+        # blocking time given, for all greenlets, no switch performed: nothing
+        with util.assert_switches(max_blocking_time=1, hub_only=False):
+            pass
+
+    def test_exception_not_supressed(self):
+
+        with self.assertRaises(NameError):
+            with util.assert_switches():
+                raise NameError()
+
+    def test_nested(self):
+        from greenlet import gettrace
+        with util.assert_switches() as outer:
+            self.assertEqual(gettrace(), outer.tracer)
+            self.assertIsNotNone(outer.tracer.active_greenlet)
+
+            with util.assert_switches() as inner:
+                self.assertEqual(gettrace(), inner.tracer)
+                self.assertEqual(inner.tracer.previous_trace_function, outer.tracer)
+
+                inner.tracer('switch', (self, self))
+
+                self.assertIs(self, inner.tracer.active_greenlet)
+                self.assertIs(self, outer.tracer.active_greenlet)
+
+            self.assertEqual(gettrace(), outer.tracer)
+
 if __name__ == '__main__':
     greentest.main()