Adding tests to check message serialization/deserialization

Avoiding python type yaml tag for message type

Signed-off-by: Emmanuel Hugonnet <[email protected]>

src/wiseagents/wise_agent_messaging.py

@@ -1,24 +1,30 @@
+import logging
 from abc import *
-from enum import Enum, auto
+from enum import StrEnum
 from typing import Callable, Optional
 
+import yaml
 from yaml import YAMLObject
+from yaml.resolver import BaseResolver
 
 
-class WiseAgentMessageType(Enum):
-    ACK = auto()
-    ALERT = auto()
-    CANNOT_ANSWER = auto()
-    QUERY = auto()
-    RESPONSE = auto()
-    ACTION_REQUEST = auto()
-    HUMAN = auto()
-    
+class WiseAgentMessageType(StrEnum):
+    ACK = "ACK"
+    ALERT = "ALERT"
+    CANNOT_ANSWER = "CANNOT_ANSWER"
+    QUERY = "QUERY"
+    RESPONSE = "RESPONSE"
+    ACTION_REQUEST = "ACTION_REQUEST"
+    HUMAN = "HUMAN"
+
 class WiseAgentEvent:
     """
     TODO
     """
 
+def wiseAgentMessageType_representer(dumper, data):
+    return dumper.represent_scalar(BaseResolver.DEFAULT_SCALAR_TAG, str(data.value))
+
 
 class WiseAgentMessage(YAMLObject):
     ''' A message that can be sent between agents. '''
@@ -47,6 +53,20 @@ def __init__(self, message: str, sender: Optional[str] = None, message_type: Opt
             self._context_name = context_name
         else:
             self._context_name = 'default'
+        self.__class__.yaml_dumper.add_representer(WiseAgentMessageType, wiseAgentMessageType_representer)
+
+    def __setstate__(self, state):
+        self._message = state["_message"]
+        self._sender =  state["_sender"]
+        self._message_type = WiseAgentMessageType(state["_message_type"])
+        self._chat_id =  state["_chat_id"]
+        self._tool_id =  state["_tool_id"]
+        self._route_response_to =  state["_route_response_to"]
+        if  state["_context_name"] is not None:
+            self._context_name = state["_context_name"]
+        else:
+            self._context_name = 'default'
+
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(message={self.message}, sender={self.sender}, message_type={self.message_type}, id={self.chat_id}, tool_id={self.tool_id}, context_name={self.context_name}, route_response_to={self.route_response_to}, route_response_to={self.route_response_to})"

test_message.yaml → tests/wiseagents/test_message.yaml

@@ -2,8 +2,7 @@
 _chat_id: '12345'
 _context_name: Weather
 _message: Hello
-_message_type: !!python/object/apply:wiseagents.wise_agent_messaging.WiseAgentMessageType
-- 1
+_message_type: ACK
 _route_response_to: Agent1
 _sender: Agent1
 _tool_id: WeatherAgent

tests/wiseagents/test_yaml_deserializer.py

@@ -4,7 +4,7 @@
 import pytest
 import yaml
 
-from wiseagents import WiseAgent, WiseAgentRegistry
+from wiseagents import WiseAgent, WiseAgentRegistry, WiseAgentMessage, WiseAgentMessageType
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -99,3 +99,24 @@ def test_assistant_desiralizer():
             deserialized_agent = yaml.load(stream, Loader=yaml.Loader)
         except yaml.YAMLError as exc:
             print(exc)
+def test_deserialize_message():
+    message = [];
+    with open(pathlib.Path().resolve() / "tests/wiseagents/test_message.yaml") as stream:
+        try:
+            message = yaml.load(stream, Loader=yaml.Loader)
+            logging.info(str(message))
+        except yaml.YAMLError as exc:
+            print(exc)
+    assert isinstance(message, WiseAgentMessage)
+    assert message.message == "Hello"
+    assert message.sender == "Agent1"
+    msgType = type(message.message_type)
+    logging.info(str(msgType))
+    assert isinstance(message.message_type, WiseAgentMessageType)
+    assert message.message_type == WiseAgentMessageType.ACK
+    assert message.chat_id == "12345"
+    assert message.tool_id == "WeatherAgent"
+    assert message.context_name =="Weather"
+    assert message.route_response_to =="Agent1"
+
+

tests/wiseagents/test_yaml_serializtion.py

@@ -1,4 +1,7 @@
 import logging
+import os
+import pathlib
+import unittest
 
 import pytest
 import yaml
@@ -147,6 +150,11 @@ def test_serialize_message():
                                tool_id="WeatherAgent", 
                                context_name="Weather", 
                                route_response_to="Agent1")
-    with open("test_message.yaml", "w") as stream:
+    with open(pathlib.Path().resolve() / "tests/wiseagents/test_serialized_message.yaml", "w") as stream:
         yaml.dump(message, stream)
+    unittest.TestCase().assertListEqual(
+        list(open(pathlib.Path().resolve() / "tests/wiseagents/test_serialized_message.yaml", "r")),
+        list(open(pathlib.Path().resolve() / "tests/wiseagents/test_message.yaml", "r")))
+    os.remove(pathlib.Path().resolve() / "tests/wiseagents/test_serialized_message.yaml")
+