Finer topic handling

osbrain/address.py

@@ -389,15 +389,15 @@ class AgentChannel():
     receiver : str
         Second AgentAddress.
     """
-    def __init__(self, kind, receiver, sender):
+    def __init__(self, kind, receiver, sender, uuid=None):
         self.kind = AgentChannelKind(kind)
         self.receiver = receiver
         self.sender = sender
         self.transport = \
             receiver.transport if receiver else sender.transport
         self.serializer = \
             receiver.serializer if receiver else sender.serializer
-        self.uuid = unique_identifier()
+        self.uuid = uuid or unique_identifier()
         # Set up pairs
         if sender:
             self.sender.channel = self
@@ -435,4 +435,5 @@ def twin(self):
         kind = self.kind.twin()
         sender = self.receiver.twin() if self.receiver is not None else None
         receiver = self.sender.twin() if self.sender is not None else None
-        return self.__class__(kind=kind, receiver=receiver, sender=sender)
+        return self.__class__(kind=kind, receiver=receiver, sender=sender,
+                              uuid=self.uuid)

osbrain/agent.py

@@ -898,6 +898,51 @@ def _connect_and_register(self, client_address, alias=None, handler=None,
         self.register(socket, register_as, alias, handler)
         return client_address
 
+    def unsubscribe_from_topic(self,
+                               socket: Union[AgentAddress, str, zmq.Socket],
+                               topic: bytes):
+        '''
+        Unsubscribe a socket from a given topic.
+
+        Parameters
+        ----------
+        socket
+            Identifier of the socket. Must be a valid entry of `self.socket`
+            for the PUB/SUB pattern and a valid entry of `self.address` for
+            the SYNC_PUB/SYNC_SUB pattern
+        topic
+            topic which we want to unsubscribe from
+        '''
+        if isinstance(self.address[socket], AgentAddress):
+            self.socket[socket].setsockopt(zmq.UNSUBSCRIBE, topic)
+        elif isinstance(self.address[socket], AgentChannel):
+            channel = self.address[socket]
+            socket = channel.receiver
+            treated_topic = channel.uuid + topic
+            self.socket[socket].setsockopt(zmq.UNSUBSCRIBE, treated_topic)
+
+    def subscribe_to_topic(self, socket: Union[AgentAddress, str, zmq.Socket],
+                           topic: bytes):
+        '''
+        Subscribe a socket to a given topic.
+
+        Parameters
+        ----------
+        socket
+            Identifier of the socket. Must be a valid entry of `self.socket`
+            for the PUB/SUB pattern and a valid entry of `self.address` for
+            the SYNC_PUB/SYNC_SUB pattern
+        topic
+            topic which we want to subscribe to
+        '''
+        if isinstance(self.address[socket], AgentAddress):
+            self.socket[socket].setsockopt(zmq.SUBSCRIBE, topic)
+        elif isinstance(self.address[socket], AgentChannel):
+            channel = self.address[socket]
+            socket = channel.receiver
+            treated_topic = channel.uuid + topic
+            self.socket[socket].setsockopt(zmq.SUBSCRIBE, treated_topic)
+
     def _handle_async_requests(self, data):
         address_uuid, uuid, response = data
         if uuid not in self._pending_requests:
@@ -932,7 +977,7 @@ def _subscribe(self, alias: str, handlers: Dict[Union[bytes, str], Any]):
         curated_handlers = topics_to_bytes(handlers)
         # Subscribe to topics
         for topic in curated_handlers.keys():
-            self.socket[alias].setsockopt(zmq.SUBSCRIBE, topic)
+            self.subscribe_to_topic(alias, topic)
         # Reset handlers
         self._set_handler(self.socket[alias], curated_handlers)
 

osbrain/tests/test_agent_sync_publications.py

@@ -73,6 +73,41 @@ def on_error(agent):
     agent.error_log.append('error')
 
 
+@pytest.mark.parametrize('socket_type', ['PUB', 'SYNC_PUB'])
+def test_change_subscription_topics_sync(nsproxy, socket_type):
+    '''
+    Test for the different options of subscribing/unsubscribing to topics
+    in the SYNC_PUB/SYNC_SUB pattern.
+    '''
+    server = run_agent('server')
+    client = run_agent('client')
+
+    addr = server.bind(socket_type, alias='pub', handler=lambda: None)
+    client.set_attr(received=[])
+    client.connect(addr, alias='sub', handler=append_received)
+
+    # Stablish a connection
+    server.each(0.1, 'send', 'pub', 'connecting...')
+    assert wait_agent_attr(client, name='received', data='connecting...')
+
+    # By default, client will be subscribed to all topics
+    server.send('pub', 'hello')
+    assert wait_agent_attr(client, name='received', data='hello')
+
+    # Only subscribe to 'TOP' topic
+    client.unsubscribe_from_topic('sub', b'')
+    client.subscribe_to_topic('sub', b'TOP')
+
+    # Message not received since 'TOP' topic not specified in the send call
+    server.send('pub', 'world')
+    assert not wait_agent_attr(client, name='received', data='world',
+                               timeout=1)
+
+    # Receive message with the topic we are subscribed to
+    server.send('pub', 'ten', topic='TOP')
+    assert wait_agent_attr(client, name='received', data='ten')
+
+
 @pytest.mark.parametrize('server', [Server, PubServer])
 def test_simple_pub_single_handler(nsproxy, server):
     """