Merge pull request 0rpc#25 from lopter/master

bombela · bombela · commit e9c57ba0377a · 2012-05-11T22:18:18.000-07:00
Add an “inspect_error” middleware hook
diff --git a/tests/test_middleware.py b/tests/test_middleware.py
@@ -524,3 +524,43 @@ def echo(self, msg):
     assert subscriber_tracer._log == [
             ('load', publisher_tracer.trace_id),
             ]
+
+def test_inspect_error_middleware():
+
+    class InspectErrorMiddleware(Tracer):
+        def __init__(self):
+            self.called = False
+            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
+
+    class Srv(object):
+        def echo(self, msg):
+            raise RuntimeError(msg)
+
+    endpoint = random_ipc_endpoint()
+
+    middleware = InspectErrorMiddleware()
+    ctx = zerorpc.Context()
+    ctx.register_middleware(middleware)
+
+    module = Srv()
+    server = zerorpc.Server(module, context=ctx)
+    server.bind(endpoint)
+    gevent.spawn(server.run)
+
+    client = zerorpc.Client()
+    client.connect(endpoint)
+
+    try:
+        client.echo('This is a test which should call the InspectErrorMiddleware')
+    except zerorpc.exceptions.RemoteError as ex:
+        assert ex.name == 'RuntimeError'
+
+    client.close()
+    server.close()
+
+    assert middleware.called is True
diff --git a/zerorpc/context.py b/zerorpc/context.py
@@ -35,12 +35,13 @@ class Context(zmq.Context):
     def __init__(self):
         self._middlewares = []
         self._middlewares_hooks = {
-                'resolve_endpoint': [],
-                'raise_error': [],
-                'call_procedure': [],
-                'load_task_context': [],
-                'get_task_context': [],
-                }
+            'resolve_endpoint': [],
+            'raise_error': [],
+            'call_procedure': [],
+            'load_task_context': [],
+            'get_task_context': [],
+            'inspect_error': []
+        }
 
     @staticmethod
     def get_instance():
@@ -71,6 +72,12 @@ def middleware_resolve_endpoint(self, endpoint):
             endpoint = functor(endpoint)
         return endpoint
 
+    def middleware_inspect_error(self, exc_type, exc_value, exc_traceback):
+        exc_info = exc_type, exc_value, exc_traceback
+        task_context = self.middleware_get_task_context()
+        for functor in self._middlewares_hooks['inspect_error']:
+            functor(task_context, exc_info)
+
     def middleware_raise_error(self, event):
         for functor in self._middlewares_hooks['raise_error']:
             functor(event)
diff --git a/zerorpc/core.py b/zerorpc/core.py
@@ -113,6 +113,9 @@ def _print_traceback(self, protocol_v1):
             traceback.print_exception(exc_type, exc_value, exc_traceback,
                     file=sys.stderr)
 
+            self._context.middleware_inspect_error(exc_type, exc_value,
+                    exc_traceback)
+
             if protocol_v1:
                 return (repr(exc_value),)
 
