From e5ff1d3c702d3d4e9d83c99e5d8226cb6248ba54 Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Sun, 2 Feb 2025 23:13:40 -0800
Subject: [PATCH] This outta coverage

---
 .../app/services/websocket-status.service.ts  |   7 +-
 src/paperless/tests/test_websockets.py        | 108 +++++++++++++++++-
 2 files changed, 109 insertions(+), 6 deletions(-)

diff --git a/src-ui/src/app/services/websocket-status.service.ts b/src-ui/src/app/services/websocket-status.service.ts
index f73b0c868..13f82412f 100644
--- a/src-ui/src/app/services/websocket-status.service.ts
+++ b/src-ui/src/app/services/websocket-status.service.ts
@@ -162,15 +162,12 @@ export class WebsocketStatusService {
       } = JSON.parse(ev.data)
 
       switch (type) {
-        case WebsocketStatusType.STATUS_UPDATE:
-          this.handleProgressUpdate(messageData as WebsocketProgressMessage)
-          break
-
         case WebsocketStatusType.DOCUMENTS_DELETED:
           this.documentDeletedSubject.next(true)
           break
 
-        default:
+        case WebsocketStatusType.STATUS_UPDATE:
+          this.handleProgressUpdate(messageData as WebsocketProgressMessage)
           break
       }
     }
diff --git a/src/paperless/tests/test_websockets.py b/src/paperless/tests/test_websockets.py
index 161bee4e6..5ba909d1c 100644
--- a/src/paperless/tests/test_websockets.py
+++ b/src/paperless/tests/test_websockets.py
@@ -5,6 +5,9 @@ from channels.testing import WebsocketCommunicator
 from django.test import TestCase
 from django.test import override_settings
 
+from documents.plugins.helpers import DocumentsStatusManager
+from documents.plugins.helpers import ProgressManager
+from documents.plugins.helpers import ProgressStatusOptions
 from paperless.asgi import application
 
 TEST_CHANNEL_LAYERS = {
@@ -22,6 +25,39 @@ class TestWebSockets(TestCase):
         self.assertFalse(connected)
         await communicator.disconnect()
 
+    @mock.patch("paperless.consumers.StatusConsumer.close")
+    @mock.patch("paperless.consumers.StatusConsumer._authenticated")
+    async def test_close_on_no_auth(self, _authenticated, mock_close):
+        _authenticated.return_value = True
+
+        communicator = WebsocketCommunicator(application, "/ws/status/")
+        connected, subprotocol = await communicator.connect()
+        self.assertTrue(connected)
+
+        message = {"type": "status_update", "data": {"task_id": "test"}}
+
+        _authenticated.return_value = False
+
+        channel_layer = get_channel_layer()
+        await channel_layer.group_send(
+            "status_updates",
+            message,
+        )
+        await communicator.receive_nothing()
+
+        mock_close.assert_called_once()
+        mock_close.reset_mock()
+
+        message = {"type": "documents_deleted", "data": {"documents": [1, 2, 3]}}
+
+        await channel_layer.group_send(
+            "status_updates",
+            message,
+        )
+        await communicator.receive_nothing()
+
+        mock_close.assert_called_once()
+
     @mock.patch("paperless.consumers.StatusConsumer._authenticated")
     async def test_auth(self, _authenticated):
         _authenticated.return_value = True
@@ -33,7 +69,7 @@ class TestWebSockets(TestCase):
         await communicator.disconnect()
 
     @mock.patch("paperless.consumers.StatusConsumer._authenticated")
-    async def test_receive(self, _authenticated):
+    async def test_receive_status_update(self, _authenticated):
         _authenticated.return_value = True
 
         communicator = WebsocketCommunicator(application, "/ws/status/")
@@ -53,3 +89,73 @@ class TestWebSockets(TestCase):
         self.assertEqual(response, message)
 
         await communicator.disconnect()
+
+    @mock.patch("paperless.consumers.StatusConsumer._authenticated")
+    async def test_receive_documents_deleted(self, _authenticated):
+        _authenticated.return_value = True
+
+        communicator = WebsocketCommunicator(application, "/ws/status/")
+        connected, subprotocol = await communicator.connect()
+        self.assertTrue(connected)
+
+        message = {"type": "documents_deleted", "data": {"documents": [1, 2, 3]}}
+
+        channel_layer = get_channel_layer()
+        await channel_layer.group_send(
+            "status_updates",
+            message,
+        )
+
+        response = await communicator.receive_json_from()
+
+        self.assertEqual(response, message)
+
+        await communicator.disconnect()
+
+    @mock.patch("channels.layers.InMemoryChannelLayer.group_send")
+    def test_manager_send_progress(self, mock_group_send):
+        with ProgressManager(task_id="test") as manager:
+            manager.send_progress(
+                ProgressStatusOptions.STARTED,
+                "Test message",
+                1,
+                10,
+                extra_args={
+                    "foo": "bar",
+                },
+            )
+
+        message = mock_group_send.call_args[0][1]
+
+        self.assertEqual(
+            message,
+            {
+                "type": "status_update",
+                "data": {
+                    "filename": None,
+                    "task_id": "test",
+                    "current_progress": 1,
+                    "max_progress": 10,
+                    "status": ProgressStatusOptions.STARTED,
+                    "message": "Test message",
+                    "foo": "bar",
+                },
+            },
+        )
+
+    @mock.patch("channels.layers.InMemoryChannelLayer.group_send")
+    def test_manager_send_documents_deleted(self, mock_group_send):
+        with DocumentsStatusManager() as manager:
+            manager.send_documents_deleted([1, 2, 3])
+
+        message = mock_group_send.call_args[0][1]
+
+        self.assertEqual(
+            message,
+            {
+                "type": "documents_deleted",
+                "data": {
+                    "documents": [1, 2, 3],
+                },
+            },
+        )