Merge pull request #21 from gridsmartercities/attribute_resolving

Attribute resolving

README.md

@@ -37,7 +37,7 @@ WSMessage is a class to represent a message to send to the websocket
 - **with_attribute**: add an attribute to the message to be sent to the websocket host
 
 ## Examples
-Testing a reponse with a body is received on connection to a websocket:
+Testing a reponse with a body is received on connection to a websocket host:
 ```py
 from pywsitest import WSTest, WSResponse
 
@@ -54,35 +54,45 @@ await ws_test.run()
 assert ws_test.is_complete()
 ```
 
-Testing a more complex set of responses and messages:
+Sending a message on connection to a websocket host:
 ```py
-from pywsitest import WSTest, WSResponse, WSMessage
+from pywsitest import WSTest, WSMessage
 
 ws_test = (
     WSTest("wss://example.com")
-    .with_parameter("Authorization", "eyJra...")
-    .with_response_timeout(15) # 15 seconds
-    .with_message_timeout(7.5)
-    .with_test_timeout(45)
     .with_message(
         WSMessage()
-        .with_attribute("type", "connect")
-        .with_attribute("body", {"chatroom": "general"})
+        .with_attribute("body", "Hello, world!")
     )
+)
+
+await ws_test.run()
+
+assert ws_test.is_complete()
+```
+
+Triggering a message to be sent when the following response is received:
+```json
+{
+    "body": {
+        "message": "Hello, world!"
+    }
+}
+```
+
+```py
+from pywsitest import WSTest, WSResponse, WSMessage
+
+ws_test = (
+    WSTest("wss://example.com")
     .with_response(
         WSResponse()
-        .with_attribute("type", "connected")
+        .with_attribute("body")
         .with_trigger(
             WSMessage()
-            .with_attribute("type", "message")
-            .with_attribute("body", "Hello, world!")
+            .with_attribute("body", "${body/message}")
         )
     )
-    .with_response(
-        WSResponse()
-        .with_attribute("type", "message")
-        .with_attribute("body", "Hello, world!")
-    )
 )
 
 await ws_test.run()

pywsitest/ws_message.py

@@ -1,4 +1,5 @@
 import json
+import re
 
 
 class WSMessage:
@@ -11,6 +12,8 @@ class WSMessage:
     Methods:
         with_attribute(key, value):
             Adds an attribute and returns the WSMessage
+        resolve(response):
+            Resolves any attributes that get their value from a parent response
 
     Usage:
         message = (
@@ -39,3 +42,31 @@ def with_attribute(self, key, value):
         """
         self.attributes[key] = value
         return self
+
+    def resolve(self, response):
+        """
+        Resolves attributes using ${path/to/property} notation with response as the source
+
+        Parameters:
+            response (dict): The response object to resolve attributes from
+
+        Returns:
+            (WSMessage): The WSMessage instance resolve was called on
+        """
+        regex = re.compile("^\$\{(.*)\}$")  # noqa: pylint - anomalous-backslash-in-string
+        for key in self.attributes:
+            value = self.attributes[key]
+            match = regex.match(str(value))
+            if match:
+                resolved = self._get_resolved_value(response, match.group(1))
+                self.attributes[key] = resolved if resolved else value
+        return self
+
+    def _get_resolved_value(self, response, path):  # noqa: pylint - no-self-use
+        resolved = response
+        for part in path.split("/"):
+            if part in resolved:
+                resolved = resolved[part]
+            else:
+                return None
+        return resolved

pywsitest/ws_test.py

@@ -196,12 +196,14 @@ async def _receive_handler(self, websocket, response):
             if expected_response.is_match(parsed_response):
                 self.received_responses.append(expected_response)
                 self.expected_responses.remove(expected_response)
-
-                for message in expected_response.triggers:
-                    await self._send_handler(websocket, message)
-
+                await self._trigger_handler(websocket, expected_response, parsed_response)
                 break
 
+    async def _trigger_handler(self, websocket, response, raw_response):
+        for message in response.triggers:
+            message = message.resolve(raw_response)
+            await self._send_handler(websocket, message)
+
     async def _send(self, websocket):
         while self.messages:
             message = self.messages.pop(0)

tests/test_ws_message.py

@@ -20,3 +20,23 @@ def test_stringify_with_attribute(self):
 
         expected_value = "{\"test\": 123}"
         self.assertEqual(expected_value, str(ws_message))
+
+    def test_resolve_attribute(self):
+        ws_message = (
+            WSMessage()
+            .with_attribute("test", 123)
+            .with_attribute("example", "${body/example}")
+        )
+
+        response = {
+            "type": "message",
+            "body": {
+                "example": 456
+            }
+        }
+
+        expected_value = "{\"test\": 123, \"example\": 456}"
+
+        ws_message = ws_message.resolve(response)
+
+        self.assertEqual(expected_value, str(ws_message))

tests/test_ws_test.py

@@ -379,3 +379,85 @@ async def test_websocket_test_receive_response_with_trigger(self, mock_ssl, mock
         self.assertTrue(ws_tester.is_complete())
         mock_socket.send.assert_called_once_with("{\"test\": 123}")
         mock_socket.close.assert_called_once()
+
+    @patch("websockets.connect")
+    @patch("ssl.SSLContext")
+    @syncify
+    async def test_websocket_test_receive_response_with_resolved_trigger(self, mock_ssl, mock_websockets):
+        ws_tester = (
+            WSTest("wss://example.com")
+            .with_response(
+                WSResponse()
+                .with_attribute("type")
+                .with_trigger(
+                    WSMessage()
+                    .with_attribute("test", "${type}")
+                )
+            )
+        )
+
+        mock_socket = MagicMock()
+        mock_socket.close = MagicMock(return_value=asyncio.Future())
+        mock_socket.close.return_value.set_result(MagicMock())
+
+        send_future = asyncio.Future()
+        send_future.set_result({})
+        mock_socket.send = MagicMock(return_value=send_future)
+
+        receive_future = asyncio.Future()
+        receive_future.set_result(json.dumps({"type": "Hello, world!"}))
+        mock_socket.recv = MagicMock(side_effect=[receive_future, asyncio.Future()])
+
+        mock_websockets.return_value = asyncio.Future()
+        mock_websockets.return_value.set_result(mock_socket)
+
+        ssl_context = MagicMock()
+        mock_ssl.return_value = ssl_context
+
+        await ws_tester.run()
+
+        self.assertTrue(ws_tester.is_complete())
+        mock_socket.send.assert_called_once_with("{\"test\": \"Hello, world!\"}")
+        mock_socket.close.assert_called_once()
+
+    @patch("websockets.connect")
+    @patch("ssl.SSLContext")
+    @syncify
+    async def test_websocket_test_receive_response_with_unresolved_trigger(self, mock_ssl, mock_websockets):
+        ws_tester = (
+            WSTest("wss://example.com")
+            .with_response(
+                WSResponse()
+                .with_attribute("type")
+                .with_trigger(
+                    WSMessage()
+                    .with_attribute("test", "${body}")
+                )
+            )
+        )
+
+        mock_socket = MagicMock()
+        mock_socket.close = MagicMock(return_value=asyncio.Future())
+        mock_socket.close.return_value.set_result(MagicMock())
+
+        send_future = asyncio.Future()
+        send_future.set_result({})
+        mock_socket.send = MagicMock(return_value=send_future)
+
+        receive_future = asyncio.Future()
+        receive_future.set_result(json.dumps({"type": "Hello, world!"}))
+        mock_socket.recv = MagicMock(
+            side_effect=[receive_future, asyncio.Future()])
+
+        mock_websockets.return_value = asyncio.Future()
+        mock_websockets.return_value.set_result(mock_socket)
+
+        ssl_context = MagicMock()
+        mock_ssl.return_value = ssl_context
+
+        await ws_tester.run()
+
+        self.assertTrue(ws_tester.is_complete())
+        mock_socket.send.assert_called_once_with(
+            "{\"test\": \"${body}\"}")
+        mock_socket.close.assert_called_once()