Add to Dispatcher the ability to forget & ignore an expected response.

git-svn-id: https://svn.erp5.org/repos/neo/trunk@2105 71dcc9de-d417-0410-9af5-da40c76e7ee4

neo/dispatcher.py

@@ -18,6 +18,7 @@
 from neo.locking import Lock
 from neo.profiling import profiler_decorator
 EMPTY = {}
+NOBODY = []
 
 def giant_lock(func):
     def wrapped(self, *args, **kw):
@@ -45,6 +46,8 @@ class Dispatcher:
         queue = self.message_table.get(id(conn), EMPTY).pop(msg_id, None)
         if queue is None:
             return False
+        elif queue is NOBODY:
+            return True
         self.queue_dict[id(queue)] -= 1
         queue.put(data)
         return True
@@ -79,6 +82,20 @@ class Dispatcher:
                 notified_set.add(queue_id)
             queue_dict[queue_id] -= 1
 
+    @giant_lock
+    @profiler_decorator
+    def forget(self, conn, msg_id):
+        """ Forget about a specific message for a specific connection.
+        Actually makes it "expected by nobody", so we know we can ignore it,
+        and not detect it as an error. """
+        message_table = self.message_table[id(conn)]
+        queue = message_table[msg_id]
+        if queue is NOBODY:
+            raise KeyError, 'Already expected by NOBODY: %r, %r' % (
+                conn, msg_id)
+        self.queue_dict[id(queue)] -= 1
+        message_table[msg_id] = NOBODY
+
     @profiler_decorator
     def registered(self, conn):
         """Check if a connection is registered into message table."""

neo/tests/testDispatcher.py

@@ -108,6 +108,21 @@ class DispatcherTests(unittest.TestCase):
         self.assertFalse(self.dispatcher.pending(queue1))
         self.assertFalse(self.dispatcher.pending(queue2))
 
+    def testForget(self):
+        conn = object()
+        queue = Queue()
+        MARKER = object()
+        # Register an expectation
+        self.dispatcher.register(conn, 1, queue)
+        # ...and forget about it
+        self.dispatcher.forget(conn, 1)
+        # If forgotten twice, it must raise a KeyError
+        self.assertRaises(KeyError, self.dispatcher.forget, conn, 1)
+        # Event arrives, return value must be True (it was expected)
+        self.assertTrue(self.dispatcher.dispatch(conn, 1, MARKER))
+        # ...but must not have reached the queue
+        self.assertTrue(queue.empty())
+
 if __name__ == '__main__':
     unittest.main()