Add test for get_handler and improve method docstring

osbrain/agent.py

@@ -547,6 +547,9 @@ def register(self, socket, address, alias=None, handler=None):
     def get_handler(self, alias):
         """
         Get the handler associated to a socket given the socket alias.
+
+        Ideally, this should only be called for alias that represent a
+        SUB socket.
         """
         return self.handler[self.socket[alias]]
 

osbrain/tests/test_agent.py

@@ -421,6 +421,21 @@ def test_invalid_handlers(nsproxy):
         agent.bind('REP', handler=1.234)
 
 
+def test_get_handler(nsproxy):
+    """
+    Make sure the actual handler is returned.
+    """
+    server = run_agent('server')
+    client = run_agent('client')
+
+    pub_addr = server.bind('PUB', alias='pub')
+    client.connect(pub_addr, alias='sub', handler=receive)
+
+    assert client.get_handler('sub')
+    with pytest.raises(KeyError):
+        server.get_handler('pub')
+
+
 def test_get_uuid_used_as_alias_for_sub_in_sync_pub_sync(nsproxy):
     """
     The function should only work for SYNC_SUB channels, and should raise