Merge pull request #1135 from NHSDigital/fix/make-relay-connections-shortlived

Use connect function as context manager

Author: steventux

manage_breast_screening/gateway/relay_service.py

@@ -8,7 +8,7 @@
 import urllib.parse
 from datetime import datetime, timezone
 
-from websockets.asyncio.client import ClientConnection, connect
+from websockets.asyncio.client import connect
 
 from .models import GatewayAction, GatewayActionStatus, Relay
 
@@ -80,9 +80,6 @@ def confirmed(self, msg: str):
 
 
 class RelayService:
-    def __init__(self):
-        self._connections: dict[str, object] = {}
-
     def send_action(self, relay: Relay, action: GatewayAction):
         """
         Synchronous wrapper around async_send_action.
@@ -97,48 +94,34 @@ async def async_send_action(
         result = SendActionResult()
 
         try:
-            conn = await self._get_connection(relay)
-            if not conn:
-                result.failed(f"No connection available for relay {relay.id}")
-                return result
-
-            await conn.send(json.dumps(action.payload))
-            result.sent(f"Sent action {action.id} to relay {relay.id}")
-
-            response = await asyncio.wait_for(
-                conn.recv(), timeout=RECEIVE_TIMEOUT_SECONDS
-            )
-            response_data = json.loads(response)
-
-            if response_data.get("status") in ("created", "processed"):
-                result.confirmed(f"Action {action.id} confirmed by gateway")
-            else:
-                result.failed(
-                    f"Unexpected response status from gateway: {response_data}"
+            relay_uri = RelayURI(relay)
+            url = relay_uri.connection_url()
+            async with connect(
+                url, compression=None, open_timeout=OPEN_CONNECTION_TIMEOUT_SECONDS
+            ) as conn:
+                await conn.send(json.dumps(action.payload))
+                result.sent(f"Sent action {action.id} to relay {relay.id}")
+
+                response = await asyncio.wait_for(
+                    conn.recv(), timeout=RECEIVE_TIMEOUT_SECONDS
                 )
+                response_data = json.loads(response)
+
+                if response_data.get("status") in ("created", "processed"):
+                    result.confirmed(f"Action {action.id} confirmed by gateway")
+                else:
+                    result.failed(
+                        f"Unexpected response status from gateway: {response_data}"
+                    )
 
         except asyncio.TimeoutError:
             result.failed(f"Timeout waiting for response from gateway {relay.id}")
 
         except Exception as e:
             result.failed(f"Error sending action to gateway {relay.id}: {e}")
-            self._connections.pop(relay.id, None)
 
         return result
 
-    async def _get_connection(self, relay: Relay) -> ClientConnection:
-        if relay.id not in self._connections:
-            self._connections[relay.id] = await self._create_connection(RelayURI(relay))
-        return self._connections[relay.id]
-
-    async def _create_connection(self, relay_uri: RelayURI) -> ClientConnection:
-        url = relay_uri.connection_url()
-        websocket = await connect(
-            url, compression=None, open_timeout=OPEN_CONNECTION_TIMEOUT_SECONDS
-        )
-        logger.info(f"Created relay connection for relay {relay_uri.relay.id}")
-        return websocket
-
     def update_gateway_action(self, action: GatewayAction, result: SendActionResult):
         """Update the GatewayAction based on the SendActionResult."""
         action.status = result.status

manage_breast_screening/gateway/tests/test_relay_service.py

@@ -81,7 +81,8 @@ def test_send_action_success(self, relay, gateway_action):
         mock_ws = AsyncMock(spec=ClientConnection)
         mock_ws.recv.return_value = json.dumps({"status": "created"})
 
-        with patch.object(subject, "_get_connection", return_value=mock_ws):
+        with patch(f"{RelayService.__module__}.connect") as mock_connect:
+            mock_connect.return_value.__aenter__.return_value = mock_ws
             subject.send_action(relay, gateway_action)
 
         mock_ws.send.assert_called_once_with(
@@ -101,7 +102,8 @@ def test_send_action_timeout(self, relay, gateway_action):
         mock_ws = AsyncMock(spec=ClientConnection)
         mock_ws.recv.side_effect = asyncio.TimeoutError
 
-        with patch.object(subject, "_get_connection", return_value=mock_ws):
+        with patch(f"{RelayService.__module__}.connect") as mock_connect:
+            mock_connect.return_value.__aenter__.return_value = mock_ws
             subject.send_action(
                 relay,
                 gateway_action,
@@ -122,7 +124,8 @@ def test_send_action_bad_response(self, relay, gateway_action):
         mock_ws = AsyncMock(spec=ClientConnection)
         mock_ws.recv.return_value = json.dumps({"unexpected": "data"})
 
-        with patch.object(subject, "_get_connection", return_value=mock_ws):
+        with patch(f"{RelayService.__module__}.connect") as mock_connect:
+            mock_connect.return_value.__aenter__.return_value = mock_ws
             subject.send_action(
                 relay,
                 gateway_action,
@@ -152,21 +155,3 @@ def test_send_action(self, relay, gateway_action):
             == "RelayService.async_send_action"
         )
         coro.close()
-
-    @pytest.mark.asyncio
-    async def test_connection_is_cached(self):
-        relay = RelayFactory.build()
-        manager = RelayService()
-
-        fake_connection = AsyncMock()
-
-        with patch.object(
-            manager,
-            "_create_connection",
-            AsyncMock(return_value=fake_connection),
-        ):
-            conn1 = await manager._get_connection(relay)
-            conn2 = await manager._get_connection(relay)
-
-        assert conn1 is conn2
-        assert relay.id in manager._connections