Use connect function as context manager

Sending a worklist item to the gateway is a brief exchange where we send a json payload and receive an acknowledgement or record a failure. The client connection we make to the gateway needn't be long lived for this simple exchange. This commit removes memoizing the connection on the RelayService class as this had no real effect. The connections would eventually close without a proper ping-pong configuration.
Using a contextmanager pattern we ensure the connection is closed when the exchange is complete or an error occurs.

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