LONK-WS: prevent duplicate subscription and implement unsubscription

chris_backend/pacsfiles/consumers.py

@@ -9,6 +9,7 @@
     validate_subscription,
     LonkWsSubscription,
     Lonk,
+    validate_unsubscription,
 )
 from pacsfiles.permissions import IsChrisOrIsPACSUserReadOnly
 
@@ -37,6 +38,9 @@ async def receive_json(self, content, **kwargs):
                 content['pacs_name'], content['SeriesInstanceUID']
             )
             return
+        if validate_unsubscription(content):
+            await self._unsubscribe_all()
+            return
         await self.close(code=400, reason='Invalid subscription')
 
     async def _subscribe(self, pacs_name: str, series_instance_uid: str):
@@ -63,6 +67,17 @@ async def _subscribe(self, pacs_name: str, series_instance_uid: str):
             await self.close(code=500)
             raise e
 
+    async def _unsubscribe_all(self):
+        """
+        Unsubscribe from *all* series notifications.
+        """
+        try:
+            await self.client.unsubscribe_all()
+            await self.send_json({'message': {'subscribed': False}})
+        except Exception as e:
+            await self.close(code=500)
+            raise e
+
     async def disconnect(self, code):
         await super().disconnect(code)
         await self.client.close()

chris_backend/pacsfiles/lonk.py

@@ -34,6 +34,14 @@ class SubscriptionRequest(TypedDict):
     action: Literal['subscribe']
 
 
+class UnsubscriptionRequest(TypedDict):
+    """
+    A request to unsubscribe from *all* series notifications.
+    """
+
+    action: Literal['unsubscribe']
+
+
 def validate_subscription(data: Any) -> TypeGuard[SubscriptionRequest]:
     if not isinstance(data, dict):
         return False
@@ -44,6 +52,12 @@ def validate_subscription(data: Any) -> TypeGuard[SubscriptionRequest]:
     )
 
 
+def validate_unsubscription(data: Any) -> TypeGuard[UnsubscriptionRequest]:
+    if not isinstance(data, dict):
+        return False
+    return data.get('action', None) == 'unsubscribe'
+
+
 class LonkProgress(TypedDict):
     """
     LONK "done" message.
@@ -114,7 +128,7 @@ class LonkClient:
 
     def __init__(self, nc: NATS):
         self._nc = nc
-        self._subscriptions: list[Subscription] = []
+        self._subscriptions: dict[str, Subscription] = {}
 
     @classmethod
     async def connect(cls, servers: str | list[str]) -> Self:
@@ -127,13 +141,23 @@ async def subscribe(
         cb: Callable[[Lonk], Awaitable[None]],
     ):
         subject = subject_of(pacs_name, series_instance_uid)
+        if (
+            subscription := self._subscriptions.get(subject, None)
+        ) is not None:
+            return subscription  # already subscribed
         cb = _curry_message2json(pacs_name, series_instance_uid, cb)
         subscription = await self._nc.subscribe(subject, cb=cb)
-        self._subscriptions.append(subscription)
+        self._subscriptions[subscription.subject] = subscription
         return subscription
 
+    async def unsubscribe_all(self):
+        await asyncio.gather(
+            *(s.unsubscribe() for s in self._subscriptions.values())
+        )
+        self._subscriptions = {}
+
     async def close(self):
-        await asyncio.gather(*(s.unsubscribe() for s in self._subscriptions))
+        await self.unsubscribe_all()
         await self._nc.close()
 
 

chris_backend/pacsfiles/tests/test_consumers.py

@@ -20,6 +20,7 @@
     LonkProgress,
     LonkDone,
     LonkError,
+    UnsubscriptionRequest,
 )
 from pacsfiles.consumers import PACSFileProgress
 from pacsfiles.tests.mocks import Mockidicom
@@ -39,15 +40,7 @@ def setUp(self):
 
     @tag('integration')
     async def test_lonk_ws(self):
-        token = await self._get_download_token()
-        app = TokenQsAuthMiddleware(PACSFileProgress.as_asgi())
-        communicator = WebsocketCommunicator(
-            app, f'v1/pacs/ws/?token={token.token}'
-        )
-        connected, subprotocol = await communicator.connect()
-        assert connected
-
-        oxidicom: Mockidicom = await Mockidicom.connect(settings.NATS_ADDRESS)
+        communicator, oxidicom = await self.connect()
 
         series1 = {'pacs_name': 'MyPACS', 'SeriesInstanceUID': '1.234.567890'}
         subscription_request = SubscriptionRequest(
@@ -103,6 +96,74 @@ async def test_lonk_ws(self):
             Lonk(message=LonkDone(done=True), **series1),
         )
 
+    @tag('integration')
+    async def test_unsubscribe(self):
+        """
+        https://chrisproject.org/docs/oxidicom/lonk-ws#unsubscribe
+        """
+        communicator, oxidicom = await self.connect()
+
+        series1 = {
+            'pacs_name': 'MyPACSUnsub',
+            'SeriesInstanceUID': '1.234.567890',
+        }
+        subscription_request = SubscriptionRequest(
+            action='subscribe', **series1
+        )
+        await communicator.send_json_to(subscription_request)
+        self.assertEqual(
+            await communicator.receive_json_from(),
+            Lonk(
+                message=LonkWsSubscription(subscribed=True),
+                **series1,
+            ),
+        )
+
+        unsubscription_request = UnsubscriptionRequest(action='unsubscribe')
+        await communicator.send_json_to(unsubscription_request)
+        self.assertEqual(
+            await communicator.receive_json_from(),
+            {'message': {'subscribed': False}},
+        )
+
+        series2 = {
+            'pacs_name': 'MyPACSUnsub',
+            'SeriesInstanceUID': '5.678.90123',
+        }
+        subscription_request = SubscriptionRequest(
+            action='subscribe', **series2
+        )
+        await communicator.send_json_to(subscription_request)
+        self.assertEqual(
+            await communicator.receive_json_from(),
+            Lonk(
+                message=LonkWsSubscription(subscribed=True),
+                **series2,
+            ),
+        )
+
+        await oxidicom.send_progress(ndicom=1, **series1)
+        await oxidicom.send_progress(ndicom=2, **series2)
+        self.assertEqual(
+            await communicator.receive_json_from(),
+            Lonk(
+                message=LonkProgress(ndicom=2),
+                **series2,  # unsubscribed from series1, should not be a message for it
+            ),
+        )
+
+    async def connect(self) -> tuple[WebsocketCommunicator, Mockidicom]:
+        token = await self._get_download_token()
+        app = TokenQsAuthMiddleware(PACSFileProgress.as_asgi())
+        communicator = WebsocketCommunicator(
+            app, f'v1/pacs/ws/?token={token.token}'
+        )
+        connected, subprotocol = await communicator.connect()
+        assert connected
+
+        oxidicom: Mockidicom = await Mockidicom.connect(settings.NATS_ADDRESS)
+        return communicator, oxidicom
+
     async def test_unauthenticated_not_connected(self):
         app = TokenQsAuthMiddleware(PACSFileProgress.as_asgi())
         communicator = WebsocketCommunicator(app, 'v1/pacs/ws/')  # no token