Return polled cmd_ids in test ethereum#50

konradkonrad · Jan 18, 2017 · 60ee530 · 60ee530
1 parent 330fa6b
commit 60ee530
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/devp2p/tests/test_discovery.py b/devp2p/tests/test_discovery.py
@@ -82,11 +82,16 @@ def receive(self, address, message):
         self.protocol.receive(address, message)
 
     def poll(self):
+        """try to receive a message
+        Returns:
+            cmd_id: the command id of the received message
+        """
         # try to receive a message
         for i, (to_address, from_address, message) in enumerate(self.messages):
             if to_address == self.address:
                 del self.messages[i]
                 self.receive(from_address, message)
+                return self.protocol.unpack(message)[1]
 
 
 def test_packing():