Merge pull request 0rpc#27 from lopter/master

bombela · bombela · commit 1a223550729c · 2012-05-21T15:55:33.000-07:00
Call the inspect_error middleware hook in Puller
diff --git a/tests/test_middleware.py b/tests/test_middleware.py
@@ -564,3 +564,45 @@ def echo(self, msg):
     server.close()
 
     assert middleware.called is True
+
+def test_inspect_error_middleware_puller():
+
+    class InspectErrorMiddleware(Tracer):
+        def __init__(self, barrier):
+            self.called = False
+            self._barrier = barrier
+            Tracer.__init__(self, identity='[server]')
+
+        def inspect_error(self, task_context, exc_info):
+            assert 'trace_id' in task_context
+            exc_type, exc_value, exc_traceback = exc_info
+            self.called = True
+            self._barrier.set()
+
+    class Srv(object):
+        def echo(self, msg):
+            raise RuntimeError(msg)
+
+    endpoint = random_ipc_endpoint()
+
+    barrier = gevent.event.Event()
+    middleware = InspectErrorMiddleware(barrier)
+    ctx = zerorpc.Context()
+    ctx.register_middleware(middleware)
+
+    module = Srv()
+    server = zerorpc.Puller(module, context=ctx)
+    server.bind(endpoint)
+    gevent.spawn(server.run)
+
+    client = zerorpc.Pusher()
+    client.connect(endpoint)
+
+    barrier.clear()
+    client.echo('This is a test which should call the InspectErrorMiddleware')
+    barrier.wait()
+
+    client.close()
+    server.close()
+
+    assert middleware.called is True
diff --git a/zerorpc/core.py b/zerorpc/core.py
@@ -320,7 +320,14 @@ def _receiver(self):
                     self._methods[event.name],
                     *event.args)
             except Exception:
-                traceback.print_exc(file=sys.stderr)
+                exc_type, exc_value, exc_traceback = sys.exc_info()
+                try:
+                    traceback.print_exception(exc_type, exc_value, exc_traceback,
+                            file=sys.stderr)
+                    self._context.middleware_inspect_error(exc_type, exc_value,
+                            exc_traceback)
+                finally:
+                    del exc_traceback
 
     def run(self):
         self._receiver_task = gevent.spawn(self._receiver)
