mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-09-15 14:18:00 +00:00
Enhancement: Make the StatusConsumer truly async (#12298)
This commit is contained in:
Trenton H
+32
-43
@@ -1,62 +1,51 @@
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from asgiref.sync import async_to_sync
|
||||
from channels.exceptions import AcceptConnection
|
||||
from channels.exceptions import DenyConnection
|
||||
from channels.generic.websocket import WebsocketConsumer
|
||||
from channels.generic.websocket import AsyncWebsocketConsumer
|
||||
|
||||
|
||||
class StatusConsumer(WebsocketConsumer):
|
||||
def _authenticated(self):
|
||||
return "user" in self.scope and self.scope["user"].is_authenticated
|
||||
class StatusConsumer(AsyncWebsocketConsumer):
|
||||
def _authenticated(self) -> bool:
|
||||
user: Any = self.scope.get("user")
|
||||
return user is not None and user.is_authenticated
|
||||
|
||||
def _can_view(self, data):
|
||||
user = self.scope.get("user") if self.scope.get("user") else None
|
||||
async def _can_view(self, data: dict[str, Any]) -> bool:
|
||||
user: Any = self.scope.get("user")
|
||||
if user is None:
|
||||
return False
|
||||
owner_id = data.get("owner_id")
|
||||
users_can_view = data.get("users_can_view", [])
|
||||
groups_can_view = data.get("groups_can_view", [])
|
||||
return (
|
||||
user.is_superuser
|
||||
or user.id == owner_id
|
||||
or user.id in users_can_view
|
||||
or any(
|
||||
user.groups.filter(pk=group_id).exists() for group_id in groups_can_view
|
||||
)
|
||||
)
|
||||
|
||||
def connect(self):
|
||||
if user.is_superuser or user.id == owner_id or user.id in users_can_view:
|
||||
return True
|
||||
|
||||
return await user.groups.filter(pk__in=groups_can_view).aexists()
|
||||
|
||||
async def connect(self) -> None:
|
||||
if not self._authenticated():
|
||||
raise DenyConnection
|
||||
else:
|
||||
async_to_sync(self.channel_layer.group_add)(
|
||||
"status_updates",
|
||||
self.channel_name,
|
||||
)
|
||||
raise AcceptConnection
|
||||
await self.close()
|
||||
return
|
||||
await self.channel_layer.group_add("status_updates", self.channel_name)
|
||||
await self.accept()
|
||||
|
||||
def disconnect(self, close_code) -> None:
|
||||
async_to_sync(self.channel_layer.group_discard)(
|
||||
"status_updates",
|
||||
self.channel_name,
|
||||
)
|
||||
async def disconnect(self, code: int) -> None:
|
||||
await self.channel_layer.group_discard("status_updates", self.channel_name)
|
||||
|
||||
def status_update(self, event) -> None:
|
||||
async def status_update(self, event: dict[str, Any]) -> None:
|
||||
if not self._authenticated():
|
||||
self.close()
|
||||
else:
|
||||
if self._can_view(event["data"]):
|
||||
self.send(json.dumps(event))
|
||||
await self.close()
|
||||
elif await self._can_view(event["data"]):
|
||||
await self.send(json.dumps(event))
|
||||
|
||||
def documents_deleted(self, event) -> None:
|
||||
async def documents_deleted(self, event: dict[str, Any]) -> None:
|
||||
if not self._authenticated():
|
||||
self.close()
|
||||
await self.close()
|
||||
else:
|
||||
self.send(json.dumps(event))
|
||||
await self.send(json.dumps(event))
|
||||
|
||||
def document_updated(self, event: Any) -> None:
|
||||
async def document_updated(self, event: dict[str, Any]) -> None:
|
||||
if not self._authenticated():
|
||||
self.close()
|
||||
else:
|
||||
if self._can_view(event["data"]):
|
||||
self.send(json.dumps(event))
|
||||
await self.close()
|
||||
elif await self._can_view(event["data"]):
|
||||
await self.send(json.dumps(event))
|
||||
|
||||
@@ -1,186 +1,175 @@
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
from channels.layers import get_channel_layer
|
||||
from channels.testing import WebsocketCommunicator
|
||||
from django.test import TestCase
|
||||
from django.test import override_settings
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
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 = {
|
||||
"default": {
|
||||
"BACKEND": "channels.layers.InMemoryChannelLayer",
|
||||
},
|
||||
}
|
||||
|
||||
class TestWebSockets:
|
||||
@pytest.fixture(autouse=True)
|
||||
def anyio_backend(self) -> str:
|
||||
return "asyncio"
|
||||
|
||||
@override_settings(CHANNEL_LAYERS=TEST_CHANNEL_LAYERS)
|
||||
class TestWebSockets(TestCase):
|
||||
@pytest.mark.anyio
|
||||
async def test_no_auth(self) -> None:
|
||||
communicator = WebsocketCommunicator(application, "/ws/status/")
|
||||
connected, _ = await communicator.connect()
|
||||
self.assertFalse(connected)
|
||||
assert not 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) -> None:
|
||||
_authenticated.return_value = True
|
||||
@pytest.mark.anyio
|
||||
async def test_close_on_no_auth(self, mocker: MockerFixture) -> None:
|
||||
mock_auth = mocker.patch(
|
||||
"paperless.consumers.StatusConsumer._authenticated",
|
||||
return_value=True,
|
||||
)
|
||||
mock_close = mocker.patch(
|
||||
"paperless.consumers.StatusConsumer.close",
|
||||
new_callable=mocker.AsyncMock,
|
||||
)
|
||||
|
||||
communicator = WebsocketCommunicator(application, "/ws/status/")
|
||||
connected, _ = await communicator.connect()
|
||||
self.assertTrue(connected)
|
||||
|
||||
message = {"type": "status_update", "data": {"task_id": "test"}}
|
||||
|
||||
_authenticated.return_value = False
|
||||
assert connected
|
||||
|
||||
mock_auth.return_value = False
|
||||
channel_layer = get_channel_layer()
|
||||
assert channel_layer is not None
|
||||
|
||||
await channel_layer.group_send(
|
||||
"status_updates",
|
||||
message,
|
||||
{"type": "status_update", "data": {"task_id": "test"}},
|
||||
)
|
||||
await communicator.receive_nothing()
|
||||
|
||||
mock_close.assert_called_once()
|
||||
mock_close.assert_awaited_once()
|
||||
mock_close.reset_mock()
|
||||
|
||||
message = {
|
||||
"type": "document_updated",
|
||||
"data": {"document_id": 10, "modified": "2026-02-17T00:00:00Z"},
|
||||
}
|
||||
|
||||
await channel_layer.group_send(
|
||||
"status_updates",
|
||||
message,
|
||||
{
|
||||
"type": "document_updated",
|
||||
"data": {"document_id": 10, "modified": "2026-02-17T00:00:00Z"},
|
||||
},
|
||||
)
|
||||
await communicator.receive_nothing()
|
||||
|
||||
mock_close.assert_called_once()
|
||||
mock_close.assert_awaited_once()
|
||||
mock_close.reset_mock()
|
||||
|
||||
message = {"type": "documents_deleted", "data": {"documents": [1, 2, 3]}}
|
||||
|
||||
await channel_layer.group_send(
|
||||
"status_updates",
|
||||
message,
|
||||
{"type": "documents_deleted", "data": {"documents": [1, 2, 3]}},
|
||||
)
|
||||
await communicator.receive_nothing()
|
||||
mock_close.assert_awaited_once()
|
||||
|
||||
mock_close.assert_called_once()
|
||||
|
||||
@mock.patch("paperless.consumers.StatusConsumer._authenticated")
|
||||
async def test_auth(self, _authenticated) -> None:
|
||||
_authenticated.return_value = True
|
||||
|
||||
communicator = WebsocketCommunicator(application, "/ws/status/")
|
||||
connected, _ = await communicator.connect()
|
||||
self.assertTrue(connected)
|
||||
|
||||
await communicator.disconnect()
|
||||
|
||||
@mock.patch("paperless.consumers.StatusConsumer._authenticated")
|
||||
async def test_receive_status_update(self, _authenticated) -> None:
|
||||
_authenticated.return_value = True
|
||||
|
||||
communicator = WebsocketCommunicator(application, "/ws/status/")
|
||||
connected, _ = await communicator.connect()
|
||||
self.assertTrue(connected)
|
||||
|
||||
message = {"type": "status_update", "data": {"task_id": "test"}}
|
||||
|
||||
channel_layer = get_channel_layer()
|
||||
await channel_layer.group_send(
|
||||
"status_updates",
|
||||
message,
|
||||
@pytest.mark.anyio
|
||||
async def test_auth(self, mocker: MockerFixture) -> None:
|
||||
mocker.patch(
|
||||
"paperless.consumers.StatusConsumer._authenticated",
|
||||
return_value=True,
|
||||
)
|
||||
|
||||
response = await communicator.receive_json_from()
|
||||
|
||||
self.assertEqual(response, message)
|
||||
communicator = WebsocketCommunicator(application, "/ws/status/")
|
||||
connected, _ = await communicator.connect()
|
||||
assert connected
|
||||
|
||||
await communicator.disconnect()
|
||||
|
||||
async def test_status_update_check_perms(self) -> None:
|
||||
@pytest.mark.anyio
|
||||
async def test_receive_status_update(self, mocker: MockerFixture) -> None:
|
||||
mocker.patch(
|
||||
"paperless.consumers.StatusConsumer._authenticated",
|
||||
return_value=True,
|
||||
)
|
||||
|
||||
communicator = WebsocketCommunicator(application, "/ws/status/")
|
||||
|
||||
communicator.scope["user"] = mock.Mock()
|
||||
communicator.scope["user"].is_authenticated = True
|
||||
communicator.scope["user"].is_superuser = False
|
||||
communicator.scope["user"].id = 1
|
||||
|
||||
connected, _ = await communicator.connect()
|
||||
self.assertTrue(connected)
|
||||
assert connected
|
||||
|
||||
# Test as owner
|
||||
message = {"type": "status_update", "data": {"task_id": "test"}}
|
||||
channel_layer = get_channel_layer()
|
||||
assert channel_layer is not None
|
||||
await channel_layer.group_send("status_updates", message)
|
||||
|
||||
assert await communicator.receive_json_from() == message
|
||||
|
||||
await communicator.disconnect()
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_status_update_check_perms(self, mocker: MockerFixture) -> None:
|
||||
user = mocker.MagicMock()
|
||||
user.is_authenticated = True
|
||||
user.is_superuser = False
|
||||
user.id = 1
|
||||
|
||||
communicator = WebsocketCommunicator(application, "/ws/status/")
|
||||
communicator.scope["user"] = user # type: ignore[typeddict-unknown-key]
|
||||
connected, _ = await communicator.connect()
|
||||
assert connected
|
||||
|
||||
channel_layer = get_channel_layer()
|
||||
assert channel_layer is not None
|
||||
|
||||
# Message received as owner
|
||||
message = {"type": "status_update", "data": {"task_id": "test", "owner_id": 1}}
|
||||
channel_layer = get_channel_layer()
|
||||
await channel_layer.group_send(
|
||||
"status_updates",
|
||||
message,
|
||||
)
|
||||
response = await communicator.receive_json_from()
|
||||
self.assertEqual(response, message)
|
||||
await channel_layer.group_send("status_updates", message)
|
||||
assert await communicator.receive_json_from() == message
|
||||
|
||||
# Test with a group that the user belongs to
|
||||
communicator.scope["user"].groups.filter.return_value.exists.return_value = True
|
||||
# Message received via group membership
|
||||
user.groups.filter.return_value.aexists = mocker.AsyncMock(return_value=True)
|
||||
message = {
|
||||
"type": "status_update",
|
||||
"data": {"task_id": "test", "owner_id": 2, "groups_can_view": [1]},
|
||||
}
|
||||
channel_layer = get_channel_layer()
|
||||
await channel_layer.group_send(
|
||||
"status_updates",
|
||||
message,
|
||||
)
|
||||
response = await communicator.receive_json_from()
|
||||
self.assertEqual(response, message)
|
||||
await channel_layer.group_send("status_updates", message)
|
||||
assert await communicator.receive_json_from() == message
|
||||
|
||||
# Test with a different owner_id
|
||||
# Message not received for different owner with no group match
|
||||
user.groups.filter.return_value.aexists = mocker.AsyncMock(return_value=False)
|
||||
message = {"type": "status_update", "data": {"task_id": "test", "owner_id": 2}}
|
||||
channel_layer = get_channel_layer()
|
||||
await channel_layer.group_send(
|
||||
"status_updates",
|
||||
message,
|
||||
)
|
||||
response = await communicator.receive_nothing()
|
||||
self.assertNotEqual(response, message)
|
||||
await channel_layer.group_send("status_updates", message)
|
||||
assert await communicator.receive_nothing()
|
||||
|
||||
await communicator.disconnect()
|
||||
|
||||
@mock.patch("paperless.consumers.StatusConsumer._authenticated")
|
||||
async def test_receive_documents_deleted(self, _authenticated) -> None:
|
||||
_authenticated.return_value = True
|
||||
@pytest.mark.anyio
|
||||
async def test_receive_documents_deleted(self, mocker: MockerFixture) -> None:
|
||||
mocker.patch(
|
||||
"paperless.consumers.StatusConsumer._authenticated",
|
||||
return_value=True,
|
||||
)
|
||||
|
||||
communicator = WebsocketCommunicator(application, "/ws/status/")
|
||||
connected, _ = await communicator.connect()
|
||||
self.assertTrue(connected)
|
||||
assert connected
|
||||
|
||||
message = {"type": "documents_deleted", "data": {"documents": [1, 2, 3]}}
|
||||
|
||||
channel_layer = get_channel_layer()
|
||||
await channel_layer.group_send(
|
||||
"status_updates",
|
||||
message,
|
||||
)
|
||||
assert channel_layer is not None
|
||||
await channel_layer.group_send("status_updates", message)
|
||||
|
||||
response = await communicator.receive_json_from()
|
||||
|
||||
self.assertEqual(response, message)
|
||||
assert await communicator.receive_json_from() == message
|
||||
|
||||
await communicator.disconnect()
|
||||
|
||||
@mock.patch("paperless.consumers.StatusConsumer._can_view")
|
||||
@mock.patch("paperless.consumers.StatusConsumer._authenticated")
|
||||
async def test_receive_document_updated(self, _authenticated, _can_view) -> None:
|
||||
_authenticated.return_value = True
|
||||
_can_view.return_value = True
|
||||
@pytest.mark.anyio
|
||||
async def test_receive_document_updated(self, mocker: MockerFixture) -> None:
|
||||
mocker.patch(
|
||||
"paperless.consumers.StatusConsumer._authenticated",
|
||||
return_value=True,
|
||||
)
|
||||
mocker.patch(
|
||||
"paperless.consumers.StatusConsumer._can_view",
|
||||
return_value=True,
|
||||
)
|
||||
|
||||
communicator = WebsocketCommunicator(application, "/ws/status/")
|
||||
connected, _ = await communicator.connect()
|
||||
self.assertTrue(connected)
|
||||
assert connected
|
||||
|
||||
message = {
|
||||
"type": "document_updated",
|
||||
@@ -192,67 +181,52 @@ class TestWebSockets(TestCase):
|
||||
"groups_can_view": [],
|
||||
},
|
||||
}
|
||||
|
||||
channel_layer = get_channel_layer()
|
||||
assert channel_layer is not None
|
||||
await channel_layer.group_send(
|
||||
"status_updates",
|
||||
message,
|
||||
)
|
||||
await channel_layer.group_send("status_updates", message)
|
||||
|
||||
response = await communicator.receive_json_from()
|
||||
|
||||
self.assertEqual(response, message)
|
||||
assert await communicator.receive_json_from() == message
|
||||
|
||||
await communicator.disconnect()
|
||||
|
||||
@mock.patch("channels.layers.InMemoryChannelLayer.group_send")
|
||||
def test_manager_send_progress(self, mock_group_send) -> None:
|
||||
def test_manager_send_progress(self, mocker: MockerFixture) -> None:
|
||||
mock_group_send = mocker.patch(
|
||||
"channels.layers.InMemoryChannelLayer.group_send",
|
||||
)
|
||||
|
||||
with ProgressManager(task_id="test") as manager:
|
||||
manager.send_progress(
|
||||
ProgressStatusOptions.STARTED,
|
||||
"Test message",
|
||||
1,
|
||||
10,
|
||||
extra_args={
|
||||
"foo": "bar",
|
||||
},
|
||||
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",
|
||||
},
|
||||
assert mock_group_send.call_args[0][1] == {
|
||||
"type": "status_update",
|
||||
"data": {
|
||||
"filename": None,
|
||||
"task_id": "test",
|
||||
"current_progress": 1,
|
||||
"max_progress": 10,
|
||||
"status": ProgressStatusOptions.STARTED,
|
||||
"message": "Test message",
|
||||
"foo": "bar",
|
||||
},
|
||||
}
|
||||
|
||||
def test_manager_send_documents_deleted(self, mocker: MockerFixture) -> None:
|
||||
mock_group_send = mocker.patch(
|
||||
"channels.layers.InMemoryChannelLayer.group_send",
|
||||
)
|
||||
|
||||
@mock.patch("channels.layers.InMemoryChannelLayer.group_send")
|
||||
def test_manager_send_documents_deleted(
|
||||
self,
|
||||
mock_group_send: mock.MagicMock,
|
||||
) -> None:
|
||||
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],
|
||||
},
|
||||
assert mock_group_send.call_args[0][1] == {
|
||||
"type": "documents_deleted",
|
||||
"data": {
|
||||
"documents": [1, 2, 3],
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user