diff --git a/zerver/lib/test_helpers.py b/zerver/lib/test_helpers.py index a62ab66cfd..7cf5a90bc4 100644 --- a/zerver/lib/test_helpers.py +++ b/zerver/lib/test_helpers.py @@ -791,3 +791,8 @@ def ratelimit_rule( with patch.dict(rules, {domain: domain_rules}), override_settings(RATE_LIMITING=True): yield + + +def consume_response(response: HttpResponseBase) -> None: + assert response.streaming + collections.deque(response, maxlen=0) diff --git a/zerver/tests/test_thumbnail.py b/zerver/tests/test_thumbnail.py index 603a206689..fbc00aef95 100644 --- a/zerver/tests/test_thumbnail.py +++ b/zerver/tests/test_thumbnail.py @@ -12,7 +12,12 @@ from django.http.request import MediaType from django.test import override_settings from zerver.lib.test_classes import ZulipTestCase -from zerver.lib.test_helpers import get_test_image_file, ratelimit_rule, read_test_image_file +from zerver.lib.test_helpers import ( + consume_response, + get_test_image_file, + ratelimit_rule, + read_test_image_file, +) from zerver.lib.thumbnail import ( BadImageError, BaseThumbnailFormat, @@ -503,11 +508,13 @@ class TestThumbnailRetrieval(ZulipTestCase): response = self.client_get(f"/user_uploads/{path_id}") self.assertEqual(response.status_code, 200) self.assertEqual(response.headers["Content-Type"], "image/gif") + consume_response(response) # Format we don't have response = self.client_get(f"/user_uploads/thumbnail/{path_id}/1x1.png") self.assertEqual(response.status_code, 404) self.assertEqual(response.headers["Content-Type"], "image/png") + consume_response(response) # Exit the block, triggering the thumbnailing worker @@ -520,6 +527,7 @@ class TestThumbnailRetrieval(ZulipTestCase): int(thumbnail_response.headers["Content-Length"]), int(response.headers["Content-Length"]), ) + consume_response(thumbnail_response) animated_response = self.client_get(f"/user_uploads/thumbnail/{path_id}/{webp_anim!s}") self.assertEqual(animated_response.status_code, 200) @@ -528,11 +536,13 @@ class TestThumbnailRetrieval(ZulipTestCase): int(thumbnail_response.headers["Content-Length"]), int(animated_response.headers["Content-Length"]), ) + consume_response(animated_response) # Invalid thumbnail format response = self.client_get(f"/user_uploads/thumbnail/{path_id}/bogus") self.assertEqual(response.status_code, 404) self.assertEqual(response.headers["Content-Type"], "image/png") + consume_response(response) # path_id for a non-image with ( @@ -546,15 +556,18 @@ class TestThumbnailRetrieval(ZulipTestCase): response = self.client_get(f"/user_uploads/thumbnail/{text_path_id}/{webp_still!s}") self.assertEqual(response.status_code, 404) self.assertEqual(response.headers["Content-Type"], "image/png") + consume_response(response) # Shrink the list of formats, and check that we can still get # the thumbnails that were generated at the time with self.mock_formats([webp_still]): response = self.client_get(f"/user_uploads/thumbnail/{path_id}/{webp_still!s}") self.assertEqual(response.status_code, 200) + consume_response(response) response = self.client_get(f"/user_uploads/thumbnail/{path_id}/{webp_anim!s}") self.assertEqual(response.status_code, 200) + consume_response(response) # Grow the format list, and check that fetching that new # format generates all of the missing formats @@ -569,6 +582,7 @@ class TestThumbnailRetrieval(ZulipTestCase): small_response = self.client_get(f"/user_uploads/thumbnail/{path_id}/{jpeg_still!s}") self.assertEqual(small_response.status_code, 200) self.assertEqual(small_response.headers["Content-Type"], "image/jpeg") + consume_response(small_response) # This made two thumbnails self.assertEqual(thumb_mock.call_count, 2) @@ -576,6 +590,7 @@ class TestThumbnailRetrieval(ZulipTestCase): big_response = self.client_get(f"/user_uploads/thumbnail/{path_id}/{big_jpeg_still!s}") self.assertEqual(big_response.status_code, 200) self.assertEqual(big_response.headers["Content-Type"], "image/jpeg") + consume_response(big_response) thumb_mock.assert_not_called() self.assertLess( @@ -598,6 +613,7 @@ class TestThumbnailRetrieval(ZulipTestCase): still_response = self.client_get(f"/user_uploads/thumbnail/{path_id}/{webp_still!s}") self.assertEqual(still_response.status_code, 200) self.assertEqual(still_response.headers["Content-Type"], "image/webp") + consume_response(still_response) # We can request -anim -- we didn't render it, but we the # "closest we rendered" logic kicks in, and we get the @@ -605,6 +621,7 @@ class TestThumbnailRetrieval(ZulipTestCase): animated_response = self.client_get(f"/user_uploads/thumbnail/{path_id}/{webp_anim!s}") self.assertEqual(animated_response.status_code, 200) self.assertEqual(animated_response.headers["Content-Type"], "image/webp") + consume_response(animated_response) # Double-check that we don't actually have the animated version, by comparing file sizes self.assertEqual( animated_response.headers["Content-Length"], @@ -614,6 +631,7 @@ class TestThumbnailRetrieval(ZulipTestCase): response = self.client_get(f"/user_uploads/thumbnail/{path_id}/{jpeg_still!s}") self.assertEqual(response.status_code, 200) self.assertEqual(response.headers["Content-Type"], "image/jpeg") + consume_response(response) def test_closest_format(self) -> None: self.login_user(self.example_user("hamlet")) diff --git a/zerver/tests/test_upload.py b/zerver/tests/test_upload.py index ed6556bad8..759d72f510 100644 --- a/zerver/tests/test_upload.py +++ b/zerver/tests/test_upload.py @@ -32,7 +32,12 @@ from zerver.lib.initial_password import initial_password from zerver.lib.realm_icon import realm_icon_url from zerver.lib.realm_logo import get_realm_logo_url from zerver.lib.test_classes import UploadSerializeMixin, ZulipTestCase -from zerver.lib.test_helpers import avatar_disk_path, get_test_image_file, ratelimit_rule +from zerver.lib.test_helpers import ( + avatar_disk_path, + consume_response, + get_test_image_file, + ratelimit_rule, +) from zerver.lib.upload import sanitize_name, upload_message_attachment from zerver.lib.upload.base import ZulipUploadBackend from zerver.lib.upload.local import LocalUploadBackend @@ -248,6 +253,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): self.logout() response = self.client_get(url) self.assertEqual(response.status_code, 200) + consume_response(response) # Deny file access since rate limited with ratelimit_rule(86400, 0, domain="spectator_attachment_access_by_file"): @@ -260,6 +266,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): with ratelimit_rule(86400, 1000, domain="spectator_attachment_access_by_file"): response = self.client_get(download_url) self.assertEqual(response.status_code, 200) + consume_response(response) with ratelimit_rule(86400, 0, domain="spectator_attachment_access_by_file"): response = self.client_get(download_url) self.assertEqual(response.status_code, 302) @@ -346,6 +353,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): ) self.assertEqual(response.status_code, 403) self.assertEqual(response.headers["Content-Type"], "image/png") + consume_response(response) def test_removed_file_download(self) -> None: """ @@ -390,6 +398,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): ) self.assertEqual(response.status_code, 404) self.assertEqual(response.headers["Content-Type"], "image/png") + consume_response(response) response = self.client_get( f"http://{hamlet.realm.host}/user_uploads/{hamlet.realm_id}/ff/gg/abc.png",