paperless-ngx/src/documents/tests/test_api_chat.py

from __future__ import annotations

from typing import TYPE_CHECKING
from unittest import mock

import pytest
from django.contrib.auth.models import Permission
from django.contrib.auth.models import User
from rest_framework import status
from rest_framework.test import APIClient
from rest_framework.test import APITestCase

if TYPE_CHECKING:
    from pytest_mock import MockerFixture


class TestChatStreamingViewInputValidation(APITestCase):
    def setUp(self) -> None:
        super().setUp()
        self.user = User.objects.create_superuser(username="temp_admin")
        self.client.force_authenticate(user=self.user)

    def _mock_ai_enabled(self) -> mock.MagicMock:
        """Return a mock AIConfig instance with ai_enabled=True."""
        m = mock.MagicMock()
        m.ai_enabled = True
        return m

    def test_oversized_question_is_rejected(self) -> None:
        with mock.patch(
            "documents.views.AIConfig",
            return_value=self._mock_ai_enabled(),
        ):
            resp = self.client.post(
                "/api/documents/chat/",
                {"q": "x" * 4001},
                format="json",
            )
        assert resp.status_code == status.HTTP_400_BAD_REQUEST

    def test_answer_is_not_compressed(self) -> None:
        """
        GIVEN:
            - A client that accepts compressed responses
        WHEN:
            - It asks the chat endpoint a question
        THEN:
            - The answer is streamed unencoded, chunk for chunk

        The stream compressors buffer, so a compressed answer arrives in one
        piece. The view cannot opt out by flagging the request: DRF's request
        wrapper proxies reads but keeps writes to itself, so the flag never
        reaches the Django request the middleware sees.
        """
        chunks = [f"token{i} " for i in range(40)]
        with (
            mock.patch(
                "documents.views.AIConfig",
                return_value=self._mock_ai_enabled(),
            ),
            mock.patch(
                "documents.views.stream_chat_with_documents",
                return_value=iter(chunks),
            ),
        ):
            resp = self.client.post(
                "/api/documents/chat/",
                {"q": "What is in my archive?"},
                format="json",
                HTTP_ACCEPT_ENCODING="gzip, deflate, br, zstd",
            )

        assert resp.status_code == status.HTTP_200_OK
        assert not resp.has_header("Content-Encoding")
        assert list(resp.streaming_content) == [c.encode() for c in chunks]

    def test_missing_question_is_rejected(self) -> None:
        with mock.patch(
            "documents.views.AIConfig",
            return_value=self._mock_ai_enabled(),
        ):
            resp = self.client.post(
                "/api/documents/chat/",
                {},
                format="json",
            )
        assert resp.status_code == status.HTTP_400_BAD_REQUEST


@pytest.mark.django_db
class TestChatStreamingViewUnrestrictedFlag:
    """The document id filter may only be skipped (``unrestricted=True``) for
    an active superuser, never for a regular user -- regardless of what
    permissions that user holds.
    """

    @pytest.fixture
    def mocked_stream_chat(self, mocker: MockerFixture) -> mock.MagicMock:
        """AI enabled, with stream_chat_with_documents patched so the view
        never touches the real vector store; returns the patched callable so
        tests can inspect how it was called.
        """
        mocker.patch("documents.views.AIConfig").return_value.ai_enabled = True
        return mocker.patch(
            "documents.views.stream_chat_with_documents",
            return_value=iter(()),
        )

    @pytest.fixture
    def viewer_client(self, user_client: APIClient, regular_user: User) -> APIClient:
        """The conftest regular-user client, granted the global
        view_document permission -- the minimum ViewDocumentsPermissions
        needs to reach the view at all. Model-level only: says nothing
        about which documents (if any) this user can actually see.
        """
        regular_user.user_permissions.add(
            *Permission.objects.filter(codename="view_document"),
        )
        return user_client

    @pytest.mark.parametrize(
        ("client_fixture", "expected_unrestricted"),
        [
            pytest.param("admin_client", True, id="superuser_is_unrestricted"),
            pytest.param("viewer_client", False, id="regular_user_is_restricted"),
        ],
    )
    def test_unrestricted_only_for_superuser(
        self,
        request: pytest.FixtureRequest,
        mocked_stream_chat: mock.MagicMock,
        client_fixture: str,
        *,
        expected_unrestricted: bool,
    ) -> None:
        """
        GIVEN:
            - A superuser, or a regular user holding the global
              view_document permission (but no object-level document access)
        WHEN:
            - They post a chat question with no document_id
        THEN:
            - stream_chat_with_documents is called with unrestricted=True
              only for the superuser; the regular user is always
              unrestricted=False, regardless of their permissions
        """
        client: APIClient = request.getfixturevalue(client_fixture)

        client.post(
            "/api/documents/chat/",
            data={"q": "What's in these documents?"},
            format="json",
        )

        assert (
            mocked_stream_chat.call_args.kwargs["unrestricted"] is expected_unrestricted
        )