Repository: opensistemas-hub/osbrain
Pull request #50 (2017-02-17 11:16:23)

A bit of code refactoring


Patch:
diff --git a/osbrain/address.py b/osbrain/address.py
index 3c8472c..ebb0ef0 100644
--- a/osbrain/address.py
+++ b/osbrain/address.py
@@ -185,11 +185,17 @@ class AgentAddressSerializer(str):
     serializer_type : str
         Serializer type (i.e.: 'raw', 'pickle'...).
     """
+    SERIALIZER_SIMPLE = ('raw', )
+    SERIALIZER_SEPARATOR = ('pickle', 'json')
+
     def __new__(cls, value):
-        if value not in ('raw', 'pickle', 'json'):
+        if value not in cls.SERIALIZER_SIMPLE + cls.SERIALIZER_SEPARATOR:
             raise ValueError('Invalid serializer type %s!' % value)
         return super().__new__(cls, value)
 
+    def __init__(self, value):
+        self.requires_separator = value in self.SERIALIZER_SEPARATOR
+
 
 class SocketAddress(object):
     """
diff --git a/osbrain/agent.py b/osbrain/agent.py
index bc5f316..2befc56 100644
--- a/osbrain/agent.py
+++ b/osbrain/agent.py
@@ -29,6 +29,9 @@ from .proxy import Proxy
 from .proxy import NSProxy
 
 
+TOPIC_SEPARATOR = b'\x80'
+
+
 def str2bytes(message):
     return message.encode('ascii')
 
@@ -37,12 +40,6 @@ def bytes2str(message):
     return message.decode('ascii')
 
 
-def check_type(x, check_type):
-    if not isinstance(x, check_type):
-        msg = 'Expected type `{}`, but got `{}`'.format(check_type, type(x))
-        raise TypeError(msg)
-
-
 def serialize_message(message, serializer):
     """
     Check if a message needs to be serialized and do it if that is the
@@ -61,21 +58,13 @@ def serialize_message(message, serializer):
         The serialized message, or the same message in case no
         serialization is needed.
     """
-    result = None
-
     if serializer == 'pickle':
-        result = pickle.dumps(message, -1)
+        return pickle.dumps(message, -1)
     if serializer == 'json':
-        result = str2bytes(json.dumps(message))
+        return str2bytes(json.dumps(message))
     if serializer == 'raw':
-        result = message
-
-    if result is None:
-        raise ValueError('Serializer not supported for serialization')
-
-    check_type(result, bytes)
-
-    return result
+        return message
+    raise ValueError('Serializer not supported for serialization')
 
 
 def deserialize_message(message, serializer):
@@ -85,7 +74,7 @@ def deserialize_message(message, serializer):
 
     Parameters
     ----------
-    message : bytes
+    message : bytes, memoryview
         The serialized message.
     serializer : AgentAddressSerializer
         The type of (de)serializer that should be used.
@@ -96,36 +85,30 @@ def deserialize_message(message, serializer):
         The deserialized message, or the same message in case no
         deserialization is needed.
     """
-    check_type(message, bytes)
-
-    result = None
-
     if serializer == 'pickle':
-        result = pickle.loads(message)
+        return pickle.loads(message)
     if serializer == 'json':
-        result = json.loads(bytes2str(message))
+        return json.loads(bytes2str(bytes(message)))
     if serializer == 'raw':
-        result = message
+        return message
+    raise ValueError('Serializer not supported for deserialization')
 
-    if result is None:
-        raise ValueError('Serializer not supported for deserialization')
 
-    return result
-
-
-def compose_message(serializer, message, topic=''):
+def compose_message(message, topic, serializer):
     """
     Compose a message and leave it ready to be sent through a socket.
 
+    This is used in PUB-SUB patterns to combine the topic and the message
+    in a single bytes buffer.
+
     Parameters
     ----------
+    message : bytes
+        Message to be composed.
+    topic : str
+        Topic to combine the message with.
     serializer : AgentAddressSerializer
         Serialization for the message part.
-    message : anything
-        Message to be serialized. The user is the one responsible for passing
-        serializable data.
-    topic : str
-        Topic.
 
     Returns
     -------
@@ -133,12 +116,10 @@ def compose_message(serializer, message, topic=''):
         The bytes representation of the final message to be sent.
     """
     assert isinstance(topic, str), 'Topic must be of `str` type!'
-    serialized = serialize_message(message=message, serializer=serializer)
     topic = str2bytes(topic)
-    if topic and serializer != 'raw':
-        separator = b'\x80'
-        return topic + separator + serialized
-    return topic + serialized
+    if serializer.requires_separator:
+        return topic + TOPIC_SEPARATOR + message
+    return topic + message
 
 
 class Agent():
@@ -865,17 +846,9 @@ class Agent():
         anything
             The content of the message passed.
         """
-        separator = b'\x80'
-
-        if serializer == 'pickle':
-            if not message.startswith(separator):
-                sepp = message.index(separator) + 1
-                message = bytes(memoryview(message)[sepp:])
-
-        if serializer == 'json':
-            if separator in message:
-                sepp = message.index(separator) + 1
-                message = bytes(memoryview(message)[sepp:])
+        if serializer.requires_separator:
+            sep = message.index(TOPIC_SEPARATOR) + 1
+            message = memoryview(message)[sep:]
 
         return deserialize_message(message=message, serializer=serializer)
 
@@ -941,7 +914,11 @@ class Agent():
         order to be sent.
         """
         serializer = self.address[address].serializer
-        message = compose_message(serializer, message, topic)
+        message = serialize_message(message=message, serializer=serializer)
+        if self.address[address].kind == 'PUB':
+            message = compose_message(message=message,
+                                      topic=topic,
+                                      serializer=serializer)
         self.socket[address].send(message)
 
     def recv(self, address):


Test patch:
diff --git a/osbrain/tests/test_address.py b/osbrain/tests/test_address.py
index 09eff51..6d81986 100644
--- a/osbrain/tests/test_address.py
+++ b/osbrain/tests/test_address.py
@@ -118,10 +118,21 @@ def test_serializer():
     """
     assert AgentAddressSerializer('raw') == 'raw'
     assert AgentAddressSerializer('pickle') == 'pickle'
+    assert AgentAddressSerializer('json') == 'json'
     with pytest.raises(ValueError):
         AgentAddressSerializer('foo')
 
 
+def test_serializer_separator():
+    """
+    Check which serializers do require a separator between topic and
+    message to be automatically inserted.
+    """
+    assert not AgentAddressSerializer('raw').requires_separator
+    assert AgentAddressSerializer('pickle').requires_separator
+    assert AgentAddressSerializer('json').requires_separator
+
+
 def test_socket_address():
     """
     Test basic SocketAddress operations: initialization and equivalence.
diff --git a/osbrain/tests/test_agent_pubsub_topics.py b/osbrain/tests/test_agent_pubsub_topics.py
index 50b35bd..813b6fc 100644
--- a/osbrain/tests/test_agent_pubsub_topics.py
+++ b/osbrain/tests/test_agent_pubsub_topics.py
@@ -1,6 +1,9 @@
 import time
 
+import pytest
+
 from osbrain import run_agent
+from osbrain.address import AgentAddressSerializer
 
 from common import nsaddr  # pragma: no flakes
 from common import nsproxy  # pragma: no flakes
@@ -10,81 +13,14 @@ def log_received_to_list(agent, message, topic=None):
     agent.received_list.append(message)
 
 
-def test_pubsub_topics_pickle(nsaddr):
-    """
-    Simple publisher-subscriber pattern test.
-
-    Different messages sent with different agents subscribed to different
-    topics are tested within this method.
-    """
-    a0 = run_agent('a0')
-    a1 = run_agent('a1')
-    a2 = run_agent('a2')
-    a3 = run_agent('a3')
-    a4 = run_agent('a4')
-    a5 = run_agent('a5')
-
-    for agent in (a1, a2, a3, a4, a5):
-        agent.set_attr(received_list=[])
-
-    addr = a0.bind('PUB', alias='pub', serializer='pickle')
-
-    a1.connect(addr, handler=log_received_to_list)
-    a2.connect(addr, handler={'foo': log_received_to_list})
-    a3.connect(addr, handler={'bar': log_received_to_list,
-                              'foo': log_received_to_list})
-    a4.connect(addr, handler={'bar': log_received_to_list})
-    a5.connect(addr, handler={'fo': log_received_to_list})
-
-    # Give some time for all the agents to connect
-    time.sleep(0.1)
-
-    # Send some messages
-    message_01 = 'Hello'
-    a0.send('pub', message_01)
-
-    message_02 = 'World'
-    a0.send('pub', message_02, topic='foo')
-
-    message_03 = 'FOO'
-    a0.send('pub', message_03, topic='foobar')
-
-    message_04 = 'BAR'
-    a0.send('pub', message_04, topic='fo')
-
-    # Give some time for all the agents to handle the message
-    time.sleep(0.1)
-
-    # Check each agent received the corresponding messages
-    assert message_01 in a1.get_attr('received_list')
-    assert message_02 in a1.get_attr('received_list')
-    assert message_03 in a1.get_attr('received_list')
-    assert message_04 in a1.get_attr('received_list')
-
-    assert message_01 not in a2.get_attr('received_list')
-    assert message_02 in a2.get_attr('received_list')
-    assert message_03 in a2.get_attr('received_list')
-    assert message_04 not in a2.get_attr('received_list')
-
-    assert message_01 not in a3.get_attr('received_list')
-    assert message_02 in a3.get_attr('received_list')
-    assert message_03 in a3.get_attr('received_list')
-    assert message_04 not in a3.get_attr('received_list')
-
-    assert message_01 not in a4.get_attr('received_list')
-    assert message_02 not in a4.get_attr('received_list')
-    assert message_03 not in a4.get_attr('received_list')
-    assert message_04 not in a4.get_attr('received_list')
-
-    assert message_01 not in a5.get_attr('received_list')
-    assert message_02 in a5.get_attr('received_list')
-    assert message_03 in a5.get_attr('received_list')
-    assert message_04 in a5.get_attr('received_list')
-
-
-def test_pubsub_topics_json(nsaddr):
+@pytest.mark.parametrize(
+    'serializer',
+    AgentAddressSerializer.SERIALIZER_SEPARATOR
+)
+def test_pubsub_topics_separator(nsaddr, serializer):
     """
-    Simple publisher-subscriber pattern test.
+    Simple publisher-subscriber pattern test using serializers that
+    require a separation in message composition.
 
     Different messages sent with different agents subscribed to different
     topics are tested within this method.
@@ -99,7 +35,7 @@ def test_pubsub_topics_json(nsaddr):
     for agent in (a1, a2, a3, a4, a5):
         agent.set_attr(received_list=[])
 
-    addr = a0.bind('PUB', alias='pub', serializer='json')
+    addr = a0.bind('PUB', alias='pub', serializer=serializer)
 
     a1.connect(addr, handler=log_received_to_list)
     a2.connect(addr, handler={'foo': log_received_to_list})
@@ -154,7 +90,11 @@ def test_pubsub_topics_json(nsaddr):
     assert message_04 in a5.get_attr('received_list')
 
 
-def test_pubsub_topics_raw(nsaddr):
+@pytest.mark.parametrize(
+    'serializer',
+    AgentAddressSerializer.SERIALIZER_SIMPLE
+)
+def test_pubsub_topics_raw(nsaddr, serializer):
     """
     Simple publisher-subscriber pattern test.
 
@@ -175,7 +115,7 @@ def test_pubsub_topics_raw(nsaddr):
     for agent in (a1, a2, a3, a4, a5):
         agent.set_attr(received_list=[])
 
-    addr = a0.bind('PUB', alias='pub', serializer='raw')
+    addr = a0.bind('PUB', alias='pub', serializer=serializer)
 
     a1.connect(addr, handler=log_received_to_list)
     a2.connect(addr, handler={'foo': log_received_to_list})
diff --git a/osbrain/tests/test_agent_serialization.py b/osbrain/tests/test_agent_serialization.py
index b838a8a..2ca8ae6 100644
--- a/osbrain/tests/test_agent_serialization.py
+++ b/osbrain/tests/test_agent_serialization.py
@@ -9,6 +9,8 @@ from osbrain import run_agent
 from osbrain.agent import serialize_message
 from osbrain.agent import deserialize_message
 from osbrain.agent import compose_message
+from osbrain.agent import TOPIC_SEPARATOR
+from osbrain.address import AgentAddressSerializer
 
 from common import nsaddr  # pragma: no flakes
 from common import nsproxy  # pragma: no flakes
@@ -19,38 +21,26 @@ def set_received(agent, message, topic=None):
 
 
 def test_message_composer():
-    msg_1 = b'Chain of bytes'
-    msg_2 = [1, 3, "Hello"]
-    topic = "test topic"
-    topic_bytes = b"test topic"
-
-    separator = b'\x80'
-
-    # Test raw serialization composing
-    assert compose_message('raw', msg_1) == b'Chain of bytes'
-    assert compose_message('raw', msg_1, topic) == b'test topicChain of bytes'
-
-    # Test pickle serialization composing
-    assert compose_message('pickle', msg_1) == pickle.dumps(msg_1, -1)
-    assert compose_message('pickle', msg_2) == pickle.dumps(msg_2, -1)
-
-    assert compose_message('pickle', msg_1, topic) \
-        == topic_bytes + separator + pickle.dumps(msg_1, -1)
-    assert compose_message('pickle', msg_2, topic) \
-        == topic_bytes + separator + pickle.dumps(msg_2, -1)
-
-    # Test json serialization composing
-    msg_1_str = 'Chain of bytes'
+    """
+    Test correct message composing.
+    """
+    message = b'message'
+    topic = 'test topic'
+    topic_bytes = b'test topic'
 
-    assert compose_message('json', msg_1_str) \
-        == json.dumps(msg_1_str).encode('ascii')
-    assert compose_message('json', msg_2) \
-        == json.dumps(msg_2).encode('ascii')
+    # Basic composing
+    for serializer in AgentAddressSerializer.SERIALIZER_SIMPLE:
+        serializer = AgentAddressSerializer(serializer)
+        assert compose_message(message, topic, serializer) \
+            == topic_bytes + message
+    for serializer in AgentAddressSerializer.SERIALIZER_SEPARATOR:
+        serializer = AgentAddressSerializer(serializer)
+        assert compose_message(message, topic, serializer) \
+            == topic_bytes + TOPIC_SEPARATOR + message
 
-    assert compose_message('json', msg_1_str, topic) \
-        == topic_bytes + separator + json.dumps(msg_1_str).encode('ascii')
-    assert compose_message('json', msg_2, topic) \
-        == topic_bytes + separator + json.dumps(msg_2).encode('ascii')
+    # Raise with wrong serializer
+    with pytest.raises(Exception):
+        compose_message(message, topic, 'foo')
 
 
 def test_serialize_message():
@@ -59,16 +49,22 @@ def test_serialize_message():
     """
     # Raw serialization
     test = b'asdf'
-    assert test == serialize_message(message=test, serializer='raw')
+    message = serialize_message(message=test, serializer='raw')
+    assert isinstance(message, bytes)
+    assert test == message
 
     # Pickle serialization
     test = [0, 1]
-    assert test == pickle.loads(serialize_message(message=test,
-                                serializer='pickle'))
+    message = serialize_message(message=test, serializer='pickle')
+    assert isinstance(message, bytes)
+    assert test == pickle.loads(message)
 
     # Json serialization
-    assert test == json.loads(serialize_message(message=test,
-                              serializer='json').decode('ascii'))
+    test = [0, 1]
+    message = serialize_message(message=test, serializer='json')
+    assert isinstance(message, bytes)
+    assert test == json.loads(message.decode('ascii'))
+
     # Un-serializable type
     with pytest.raises(TypeError):
         serialize_message(message=b'Hello', serializer='json')
@@ -96,14 +92,6 @@ def test_deserialize_message():
         deserialize_message(message=json.dumps(test).encode('ascii'),
                             serializer='json')
 
-    # Wrong type
-    with pytest.raises(TypeError):
-        deserialize_message(message='Hello', serializer='raw')
-    with pytest.raises(TypeError):
-        deserialize_message(message=[1, 2, 3], serializer='pickle')
-    with pytest.raises(TypeError):
-        deserialize_message(message=None, serializer='json')
-
     # Incorrect serializer
     with pytest.raises(ValueError):
         deserialize_message(message=b'x', serializer='foo')