paperless-ngx/src/paperless/tests/test_compression_middleware.py

from django.http import HttpResponse
from django.http import StreamingHttpResponse
from django.test import RequestFactory
from django.test import TestCase

from paperless.middleware import StreamAwareCompressionMiddleware


class TestStreamAwareCompressionMiddleware(TestCase):
    def setUp(self) -> None:
        super().setUp()
        self.factory = RequestFactory()
        self.middleware = StreamAwareCompressionMiddleware(lambda request: None)

    def _request(self):
        return self.factory.get(
            "/api/documents/chat/",
            HTTP_ACCEPT_ENCODING="gzip, deflate, br, zstd",
        )

    def test_event_stream_is_not_compressed(self) -> None:
        """
        GIVEN:
            - A server-sent event response produced chunk by chunk
        WHEN:
            - The compression middleware processes it
        THEN:
            - It is passed through unencoded, one wire chunk per source chunk
        """
        chunks = [f"token{i} ".encode() for i in range(40)]
        response = StreamingHttpResponse(
            iter(chunks),
            content_type="text/event-stream",
        )

        response = self.middleware.process_response(self._request(), response)

        assert not response.has_header("Content-Encoding")
        assert list(response.streaming_content) == chunks

    def test_regular_response_is_still_compressed(self) -> None:
        """
        GIVEN:
            - An ordinary response large enough to be worth compressing
        WHEN:
            - The compression middleware processes it
        THEN:
            - It is compressed as before
        """
        response = HttpResponse(b"a" * 5000, content_type="application/json")

        response = self.middleware.process_response(self._request(), response)

        assert response.has_header("Content-Encoding")