Add tests

Signed-off-by: Nicholas Rempel <[email protected]>

aries_cloudagent/core/dispatcher.py

@@ -210,6 +210,7 @@ async def make_message(self, parsed_msg: dict) -> AgentMessage:
 
         registry: ProtocolRegistry = await self.context.inject(ProtocolRegistry)
         message_type = parsed_msg.get("@type")
+
         if not message_type:
             raise MessageParseError("Message does not contain '@type' parameter")
 

aries_cloudagent/core/protocol_registry.py

@@ -136,7 +136,7 @@ def resolve_message_class(self, message_type: str) -> type:
             return ClassLoader.load_class(msg_cls)
 
         # Support registered modules (not path as string)
-        if msg_cls:
+        elif msg_cls:
             return msg_cls
 
         # Try and route via min/maj version matching
@@ -165,7 +165,10 @@ def resolve_message_class(self, message_type: str) -> type:
                             + f" Received {parsed_type_string['minor_version']}."
                         )
 
-                    return ClassLoader.load_class(proto["message_module"])
+                    if isinstance(proto["message_module"], str):
+                        return ClassLoader.load_class(msg_cls)
+                    elif proto["message_module"]:
+                        return proto["message_module"]
 
         return None
 

aries_cloudagent/core/tests/test_dispatcher.py

@@ -50,7 +50,7 @@ class StubAgentMessage(AgentMessage):
     class Meta:
         handler_class = "StubAgentMessageHandler"
         schema_class = "StubAgentMessageSchema"
-        message_type = "stub-message"
+        message_type = "proto-name/1.1/message-type"
 
 
 class StubAgentMessageSchema(AgentMessageSchema):
@@ -63,6 +63,22 @@ async def handle(self, context, responder):
         pass
 
 
+class StubV1_2AgentMessage(AgentMessage):
+    class Meta:
+        handler_class = "StubV1_2AgentMessageHandler"
+        schema_class = "StubV1_2AgentMessageSchema"
+        message_type = "proto-name/1.2/message-type"
+
+class StubV1_2AgentMessageSchema(AgentMessageSchema):
+    class Meta:
+        model_class = StubV1_2AgentMessage
+
+
+class StubV1_2AgentMessageHandler:
+    async def handle(self, context, responder):
+        pass
+
+
 class TestDispatcher(AsyncTestCase):
     async def test_dispatch(self):
         context = make_context()
@@ -87,6 +103,92 @@ async def test_dispatch(self):
                 handler_mock.call_args[0][2], test_module.DispatcherResponder
             )
 
+    async def test_dispatch_versioned_message(self):
+        context = make_context()
+        context.enforce_typing = False
+        registry = await context.inject(ProtocolRegistry)
+        registry.register_message_types(
+            {StubAgentMessage.Meta.message_type: StubAgentMessage},
+            version_definition={
+                "major_version": 1,
+                "minimum_minor_version": 0,
+                "current_minor_version": 1,
+                "path": "v1_1",
+            },
+        )
+        dispatcher = test_module.Dispatcher(context)
+        await dispatcher.setup()
+        rcv = Receiver()
+        message = {"@type": StubAgentMessage.Meta.message_type}
+
+        with async_mock.patch.object(
+            StubAgentMessageHandler, "handle", autospec=True
+        ) as handler_mock:
+            await dispatcher.queue_message(make_inbound(message), rcv.send)
+            await dispatcher.task_queue
+            handler_mock.assert_awaited_once()
+            assert isinstance(handler_mock.call_args[0][1].message, StubAgentMessage)
+            assert isinstance(
+                handler_mock.call_args[0][2], test_module.DispatcherResponder
+            )
+
+    async def test_dispatch_versioned_message_handle_greater_succeeds(self):
+        context = make_context()
+        context.enforce_typing = False
+        registry = await context.inject(ProtocolRegistry)
+        registry.register_message_types(
+            {StubAgentMessage.Meta.message_type: StubAgentMessage},
+            version_definition={
+                "major_version": 1,
+                "minimum_minor_version": 0,
+                "current_minor_version": 1,
+                "path": "v1_1",
+            },
+        )
+        dispatcher = test_module.Dispatcher(context)
+        await dispatcher.setup()
+        rcv = Receiver()
+        message = {"@type": StubV1_2AgentMessage.Meta.message_type}
+
+        with async_mock.patch.object(
+            StubAgentMessageHandler, "handle", autospec=True
+        ) as handler_mock:
+            await dispatcher.queue_message(make_inbound(message), rcv.send)
+            await dispatcher.task_queue
+            handler_mock.assert_awaited_once()
+            assert isinstance(handler_mock.call_args[0][1].message, StubAgentMessage)
+            assert isinstance(
+                handler_mock.call_args[0][2], test_module.DispatcherResponder
+            )
+
+    async def test_dispatch_versioned_message_fail(self):
+        context = make_context()
+        context.enforce_typing = False
+        registry = await context.inject(ProtocolRegistry)
+        registry.register_message_types(
+            {StubV1_2AgentMessage.Meta.message_type: StubV1_2AgentMessage},
+            version_definition={
+                "major_version": 1,
+                "minimum_minor_version": 2,
+                "current_minor_version": 2,
+                "path": "v1_2",
+            },
+        )
+        dispatcher = test_module.Dispatcher(context)
+        await dispatcher.setup()
+        rcv = Receiver()
+        message = {"@type": StubAgentMessage.Meta.message_type}
+
+        with async_mock.patch.object(
+            StubAgentMessageHandler, "handle", autospec=True
+        ) as handler_mock:
+            await dispatcher.queue_message(make_inbound(message), rcv.send)
+            await dispatcher.task_queue
+            assert rcv.messages and isinstance(rcv.messages[0][1], OutboundMessage)
+            payload = json.loads(rcv.messages[0][1].payload)
+            assert payload["@type"] == ProblemReport.Meta.message_type
+
+
     async def test_bad_message_dispatch(self):
         dispatcher = test_module.Dispatcher(make_context())
         await dispatcher.setup()