Added the call_procedure middleware

aluzzardi · bombela · commit 53e8e6595767 · 2012-03-12T17:27:35.000-07:00
diff --git a/tests/test_middleware.py b/tests/test_middleware.py
@@ -152,3 +152,100 @@ def raise_error(self, event):
 
     client.close()
     srv.close()
+
+
+def test_call_procedure():
+    c = zerorpc.Context()
+
+    def test(argument):
+        return 'ret_real:' + argument
+    assert c.middleware_call_procedure(test, 'dummy') == 'ret_real:dummy'
+
+    def middleware_1(procedure, *args, **kwargs):
+        return 'ret_middleware_1:' + procedure(*args, **kwargs)
+    cnt = c.register_middleware({
+        'call_procedure': middleware_1
+        })
+    assert cnt == 1
+    assert c.middleware_call_procedure(test, 'dummy') == \
+        'ret_middleware_1:ret_real:dummy'
+
+    def middleware_2(procedure, *args, **kwargs):
+        return 'ret_middleware_2:' + procedure(*args, **kwargs)
+    cnt = c.register_middleware({
+        'call_procedure': middleware_2
+        })
+    assert cnt == 1
+    assert c.middleware_call_procedure(test, 'dummy') == \
+        'ret_middleware_2:ret_middleware_1:ret_real:dummy'
+
+    def mangle_arguments(procedure, *args, **kwargs):
+        return procedure(args[0].upper())
+    cnt = c.register_middleware({
+        'call_procedure': mangle_arguments
+        })
+    assert cnt == 1
+    assert c.middleware_call_procedure(test, 'dummy') == \
+        'ret_middleware_2:ret_middleware_1:ret_real:DUMMY'
+
+    endpoint = 'ipc://test_call_procedure'
+
+    # client/server
+    class Server(zerorpc.Server):
+        def test(self, argument):
+            return 'ret_real:' + argument
+    server = Server(heartbeat=1, context=c)
+    server.bind(endpoint)
+    gevent.spawn(server.run)
+    client = zerorpc.Client(heartbeat=1, context=c)
+    client.connect(endpoint)
+    assert client.test('dummy') == \
+        'ret_middleware_2:ret_middleware_1:ret_real:DUMMY'
+    client.close()
+    server.close()
+
+    # push/pull
+    trigger = gevent.event.Event()
+    class Puller(zerorpc.Puller):
+        argument = None
+
+        def test(self, argument):
+            self.argument = argument
+            trigger.set()
+            return self.argument
+
+    puller = Puller(context=c)
+    puller.bind(endpoint)
+    gevent.spawn(puller.run)
+    pusher = zerorpc.Pusher(context=c)
+    pusher.connect(endpoint)
+    trigger.clear()
+    pusher.test('dummy')
+    trigger.wait()
+    assert puller.argument == 'DUMMY'
+    #FIXME: These seems to be broken
+    # pusher.close()
+    # puller.close()
+
+    # pub/sub
+    trigger = gevent.event.Event()
+    class Subscriber(zerorpc.Subscriber):
+        argument = None
+
+        def test(self, argument):
+            self.argument = argument
+            trigger.set()
+            return self.argument
+
+    subscriber = Subscriber(context=c)
+    subscriber.bind(endpoint)
+    gevent.spawn(subscriber.run)
+    publisher = zerorpc.Publisher(context=c)
+    publisher.connect(endpoint)
+    trigger.clear()
+    publisher.test('dummy')
+    trigger.wait()
+    assert subscriber.argument == 'DUMMY'
+    #FIXME: These seems to be broken
+    # publisher.close()
+    # subscriber.close()
diff --git a/zerorpc/context.py b/zerorpc/context.py
@@ -35,7 +35,8 @@ def __init__(self):
         self._middlewares = []
         self._middlewares_hooks = {
                 'resolve_endpoint': [],
-                'raise_error': []
+                'raise_error': [],
+                'call_procedure': []
                 }
 
     @staticmethod
@@ -70,3 +71,16 @@ def middleware_resolve_endpoint(self, endpoint):
     def middleware_raise_error(self, event):
         for functor in self._middlewares_hooks['raise_error']:
             functor(event)
+
+    def middleware_call_procedure(self, procedure, *args, **kwargs):
+        class chain(object):
+            def __init__(self, fct, next):
+                self.fct = fct
+                self.next = next
+
+            def __call__(self, *args, **kwargs):
+                return self.fct(self.next, *args, **kwargs)
+
+        for functor in self._middlewares_hooks['call_procedure']:
+            procedure = chain(functor, procedure)
+        return procedure(*args, **kwargs)
diff --git a/zerorpc/core.py b/zerorpc/core.py
@@ -70,15 +70,15 @@ def _zerorpc_args(self):
 
 class PatternReqRep():
 
-    def process_call(self, socket, event, functor):
-        result = functor(*event.args)
+    def process_call(self, context, socket, event, functor):
+        result = context.middleware_call_procedure(functor, *event.args)
         socket.emit('OK', (result,))
 
     def accept_answer(self, event):
         return True
 
-    def process_answer(self, socket, event, method, timeout,
-        raise_remote_error):
+    def process_answer(self, context, socket, event, method, timeout,
+            raise_remote_error):
         result = event.args[0]
         if event.name == 'ERR':
             raise_remote_error(event)
@@ -93,15 +93,16 @@ class rep(DecoratorBase):
 
 class PatternReqStream():
 
-    def process_call(self, socket, event, functor):
-        for result in iter(functor(*event.args)):
+    def process_call(self, context, socket, event, functor):
+        for result in iter(context.middleware_call_procedure(functor,
+                *event.args)):
             socket.emit('STREAM', result)
         socket.emit('STREAM_DONE', None)
 
     def accept_answer(self, event):
         return event.name in ('STREAM', 'STREAM_DONE')
 
-    def process_answer(self, socket, event, method, timeout,
+    def process_answer(self, context, socket, event, method, timeout,
             raise_remote_error):
         def iterator(event):
             while event.name == 'STREAM':
@@ -197,7 +198,7 @@ def _async_task(self, initial_event):
             functor = self._methods.get(event.name, None)
             if functor is None:
                 raise NameError(event.name)
-            functor.pattern.process_call(socket, event, functor)
+            functor.pattern.process_call(self._context, socket, event, functor)
         except Exception:
             try:
                 exc_type, exc_value, exc_traceback = sys.exc_info()
@@ -282,8 +283,8 @@ def _process_response(self, method, socket, timeout):
                         'calling remote method {0}'.format(method))
 
             pattern = self._select_pattern(event)
-            return pattern.process_answer(socket, event, method, timeout,
-                    self._raise_remote_error)
+            return pattern.process_answer(self._context, socket, event, method,
+                timeout, self._raise_remote_error)
         except:
             socket.close()
             socket.channel.close()
@@ -365,7 +366,9 @@ def _receiver(self):
             try:
                 if event.name not in self._methods:
                     raise NameError(event.name)
-                self._methods[event.name](*event.args)
+                self._context.middleware_call_procedure(
+                    self._methods[event.name],
+                    *event.args)
             except Exception:
                 traceback.print_exc(file=sys.stderr)
 
