From 92c83c1df489975f8fe887f190e5d0af7c9534a2 Mon Sep 17 00:00:00 2001 From: Anders Kaseorg Date: Thu, 15 Jun 2023 16:07:40 -0700 Subject: [PATCH] tests: Remove assert_streaming_content helper in favor of getvalue. Signed-off-by: Anders Kaseorg --- zerver/lib/test_classes.py | 6 ----- zerver/tests/test_auth_backends.py | 4 +-- zerver/tests/test_realm_export.py | 2 +- zerver/tests/test_upload.py | 40 ++++++++++++++---------------- zerver/tests/test_upload_local.py | 4 +-- 5 files changed, 22 insertions(+), 34 deletions(-) diff --git a/zerver/lib/test_classes.py b/zerver/lib/test_classes.py index 8857444216..25df47e351 100644 --- a/zerver/lib/test_classes.py +++ b/zerver/lib/test_classes.py @@ -36,7 +36,6 @@ from django.db.migrations.executor import MigrationExecutor from django.db.migrations.state import StateApps from django.db.utils import IntegrityError from django.http import HttpRequest, HttpResponse -from django.http.response import StreamingHttpResponse from django.test import TestCase from django.test.client import BOUNDARY, MULTIPART_CONTENT, encode_multipart from django.test.testcases import SerializeMixin @@ -1111,11 +1110,6 @@ Output: return [subscription.user_profile for subscription in subscriptions] - def assert_streaming_content(self, response: "TestHttpResponse", result: bytes) -> None: - assert isinstance(response, StreamingHttpResponse) - data = b"".join(response.streaming_content) - self.assertEqual(result, data) - def assert_json_success( self, result: Union["TestHttpResponse", HttpResponse], diff --git a/zerver/tests/test_auth_backends.py b/zerver/tests/test_auth_backends.py index fe61c5a616..f8525b26cc 100644 --- a/zerver/tests/test_auth_backends.py +++ b/zerver/tests/test_auth_backends.py @@ -6307,8 +6307,8 @@ class TestLDAP(ZulipLDAPTestCase): "rb", ) as f: example_avatar = f.read() - self.assert_streaming_content( - response, resize_avatar(example_avatar, DEFAULT_AVATAR_SIZE) + self.assertEqual( + response.getvalue(), resize_avatar(example_avatar, DEFAULT_AVATAR_SIZE) ) @override_settings(AUTHENTICATION_BACKENDS=("zproject.backends.ZulipLDAPAuthBackend",)) diff --git a/zerver/tests/test_realm_export.py b/zerver/tests/test_realm_export.py index 783f9e0804..4e5cfa40b2 100644 --- a/zerver/tests/test_realm_export.py +++ b/zerver/tests/test_realm_export.py @@ -176,7 +176,7 @@ class RealmExportTest(ZulipTestCase): export_path = orjson.loads(extra_data).get("export_path") response = self.client_get(export_path) self.assertEqual(response.status_code, 200) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") result = self.client_get("/json/export/realm") response_dict = self.assert_json_success(result) diff --git a/zerver/tests/test_upload.py b/zerver/tests/test_upload.py index 13e53289c9..7939282f68 100644 --- a/zerver/tests/test_upload.py +++ b/zerver/tests/test_upload.py @@ -10,7 +10,6 @@ from unittest.mock import patch import orjson from django.conf import settings -from django.http.response import StreamingHttpResponse from django.utils.timezone import now as timezone_now from PIL import Image from urllib3 import encode_multipart_formdata @@ -81,11 +80,11 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): self.logout() response = self.api_get(self.example_user("hamlet"), url) self.assertEqual(response.status_code, 200) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") # Files uploaded through the API should be accessible via the web client self.login("hamlet") - self.assert_streaming_content(self.client_get(url), b"zulip!") + self.assertEqual(self.client_get(url).getvalue(), b"zulip!") def test_mobile_api_endpoint(self) -> None: """ @@ -114,7 +113,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): response = self.client_get(url, {"api_key": get_api_key(user_profile)}) self.assertEqual(response.status_code, 200) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") def test_file_too_big_failure(self) -> None: """ @@ -194,12 +193,12 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): # In the future, local file requests will follow the same style as S3 # requests; they will be first authenticated and redirected - self.assert_streaming_content(self.client_get(url), b"zulip!") + self.assertEqual(self.client_get(url).getvalue(), b"zulip!") # Check the download endpoint download_url = url.replace("/user_uploads/", "/user_uploads/download/") result = self.client_get(download_url) - self.assert_streaming_content(result, b"zulip!") + self.assertEqual(result.getvalue(), b"zulip!") self.assertIn("attachment;", result.headers["Content-Disposition"]) # check if DB has attachment marked as unclaimed @@ -223,7 +222,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): # The generated URL has a token authorizing the requestor to access the file # without being logged in. self.logout() - self.assert_streaming_content(self.client_get(url_only_url), b"zulip!") + self.assertEqual(self.client_get(url_only_url).getvalue(), b"zulip!") # The original url shouldn't work when logged out: result = self.client_get(url) self.assertEqual(result.status_code, 403) @@ -314,7 +313,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): url_only_url = data["url"] self.logout() - self.assert_streaming_content(self.client_get(url_only_url), b"zulip!") + self.assertEqual(self.client_get(url_only_url).getvalue(), b"zulip!") # After over 60 seconds, the token should become invalid: with mock.patch("django.core.signing.time.time", return_value=start_time + 61): @@ -741,7 +740,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): self.login_user(user_1) response = self.client_get(url, subdomain=test_subdomain) self.assertEqual(response.status_code, 200) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") self.logout() # Confirm other cross-realm users can't read it. @@ -784,7 +783,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): with self.assert_database_query_count(5): response = self.client_get(url) self.assertEqual(response.status_code, 200) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") self.logout() # Subscribed user who received the message should be able to view file @@ -792,7 +791,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): with self.assert_database_query_count(6): response = self.client_get(url) self.assertEqual(response.status_code, 200) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") self.logout() def assert_cannot_access_file(user: UserProfile) -> None: @@ -845,7 +844,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): with self.assert_database_query_count(5): response = self.client_get(url) self.assertEqual(response.status_code, 200) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") self.logout() # Originally subscribed user should be able to view file @@ -853,7 +852,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): with self.assert_database_query_count(6): response = self.client_get(url) self.assertEqual(response.status_code, 200) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") self.logout() # Subscribed user who did not receive the message should also be able to view file @@ -861,7 +860,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): with self.assert_database_query_count(9): response = self.client_get(url) self.assertEqual(response.status_code, 200) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") self.logout() # It takes a few extra queries to verify access because of shared history. @@ -921,7 +920,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): with self.assert_database_query_count(9): response = self.client_get(url) self.assertEqual(response.status_code, 200) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") with self.assert_database_query_count(6): self.assertTrue(validate_attachment_request(user, fp_path_id)) @@ -949,7 +948,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase): for user in subscribed_users + unsubscribed_users: self.login_user(user) response = self.client_get(url) - self.assert_streaming_content(response, b"zulip!") + self.assertEqual(response.getvalue(), b"zulip!") self.logout() def test_serve_local(self) -> None: @@ -1309,8 +1308,7 @@ class AvatarTest(UploadSerializeMixin, ZulipTestCase): if rfname is not None: response = self.client_get(url) - assert isinstance(response, StreamingHttpResponse) - data = b"".join(response.streaming_content) + data = response.getvalue() self.assertEqual(Image.open(io.BytesIO(data)).size, (100, 100)) # Verify that the medium-size avatar was created @@ -1590,8 +1588,7 @@ class RealmIconTest(UploadSerializeMixin, ZulipTestCase): if rfname is not None: response = self.client_get(url) - assert isinstance(response, StreamingHttpResponse) - data = b"".join(response.streaming_content) + data = response.getvalue() self.assertEqual(Image.open(io.BytesIO(data)).size, (100, 100)) def test_invalid_icons(self) -> None: @@ -1773,8 +1770,7 @@ class RealmLogoTest(UploadSerializeMixin, ZulipTestCase): if rfname is not None: response = self.client_get(logo_url) - assert isinstance(response, StreamingHttpResponse) - data = b"".join(response.streaming_content) + data = response.getvalue() # size should be 100 x 100 because thumbnail keeps aspect ratio # while trying to fit in a 800 x 100 box without losing part of the image self.assertEqual(Image.open(io.BytesIO(data)).size, (100, 100)) diff --git a/zerver/tests/test_upload_local.py b/zerver/tests/test_upload_local.py index f4947c898e..034ea92322 100644 --- a/zerver/tests/test_upload_local.py +++ b/zerver/tests/test_upload_local.py @@ -5,7 +5,6 @@ from io import BytesIO, StringIO from urllib.parse import urlparse from django.conf import settings -from django.http.response import StreamingHttpResponse from PIL import Image import zerver.lib.upload @@ -150,8 +149,7 @@ class LocalStorageTest(UploadSerializeMixin, ZulipTestCase): # We get a resized avatar from it image_data = read_test_image_file("img.png") resized_avatar = resize_avatar(image_data) - assert isinstance(result, StreamingHttpResponse) - self.assertEqual(resized_avatar, b"".join(result.streaming_content)) + self.assertEqual(resized_avatar, result.getvalue()) with self.settings(DEVELOPMENT=False): # In production, this is an X-Accel-Redirect to the