tests: Remove assert_streaming_content helper in favor of getvalue.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
This commit is contained in:
Anders Kaseorg 2023-06-15 16:07:40 -07:00 committed by Anders Kaseorg
parent 5901ffb0ab
commit 92c83c1df4
5 changed files with 22 additions and 34 deletions

View File

@ -36,7 +36,6 @@ from django.db.migrations.executor import MigrationExecutor
from django.db.migrations.state import StateApps from django.db.migrations.state import StateApps
from django.db.utils import IntegrityError from django.db.utils import IntegrityError
from django.http import HttpRequest, HttpResponse from django.http import HttpRequest, HttpResponse
from django.http.response import StreamingHttpResponse
from django.test import TestCase from django.test import TestCase
from django.test.client import BOUNDARY, MULTIPART_CONTENT, encode_multipart from django.test.client import BOUNDARY, MULTIPART_CONTENT, encode_multipart
from django.test.testcases import SerializeMixin from django.test.testcases import SerializeMixin
@ -1111,11 +1110,6 @@ Output:
return [subscription.user_profile for subscription in subscriptions] 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( def assert_json_success(
self, self,
result: Union["TestHttpResponse", HttpResponse], result: Union["TestHttpResponse", HttpResponse],

View File

@ -6307,8 +6307,8 @@ class TestLDAP(ZulipLDAPTestCase):
"rb", "rb",
) as f: ) as f:
example_avatar = f.read() example_avatar = f.read()
self.assert_streaming_content( self.assertEqual(
response, resize_avatar(example_avatar, DEFAULT_AVATAR_SIZE) response.getvalue(), resize_avatar(example_avatar, DEFAULT_AVATAR_SIZE)
) )
@override_settings(AUTHENTICATION_BACKENDS=("zproject.backends.ZulipLDAPAuthBackend",)) @override_settings(AUTHENTICATION_BACKENDS=("zproject.backends.ZulipLDAPAuthBackend",))

View File

@ -176,7 +176,7 @@ class RealmExportTest(ZulipTestCase):
export_path = orjson.loads(extra_data).get("export_path") export_path = orjson.loads(extra_data).get("export_path")
response = self.client_get(export_path) response = self.client_get(export_path)
self.assertEqual(response.status_code, 200) 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") result = self.client_get("/json/export/realm")
response_dict = self.assert_json_success(result) response_dict = self.assert_json_success(result)

View File

@ -10,7 +10,6 @@ from unittest.mock import patch
import orjson import orjson
from django.conf import settings from django.conf import settings
from django.http.response import StreamingHttpResponse
from django.utils.timezone import now as timezone_now from django.utils.timezone import now as timezone_now
from PIL import Image from PIL import Image
from urllib3 import encode_multipart_formdata from urllib3 import encode_multipart_formdata
@ -81,11 +80,11 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
self.logout() self.logout()
response = self.api_get(self.example_user("hamlet"), url) response = self.api_get(self.example_user("hamlet"), url)
self.assertEqual(response.status_code, 200) 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 # Files uploaded through the API should be accessible via the web client
self.login("hamlet") 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: 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)}) response = self.client_get(url, {"api_key": get_api_key(user_profile)})
self.assertEqual(response.status_code, 200) 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: 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 # In the future, local file requests will follow the same style as S3
# requests; they will be first authenticated and redirected # 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 # Check the download endpoint
download_url = url.replace("/user_uploads/", "/user_uploads/download/") download_url = url.replace("/user_uploads/", "/user_uploads/download/")
result = self.client_get(download_url) 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"]) self.assertIn("attachment;", result.headers["Content-Disposition"])
# check if DB has attachment marked as unclaimed # 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 # The generated URL has a token authorizing the requestor to access the file
# without being logged in. # without being logged in.
self.logout() 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: # The original url shouldn't work when logged out:
result = self.client_get(url) result = self.client_get(url)
self.assertEqual(result.status_code, 403) self.assertEqual(result.status_code, 403)
@ -314,7 +313,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
url_only_url = data["url"] url_only_url = data["url"]
self.logout() 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: # After over 60 seconds, the token should become invalid:
with mock.patch("django.core.signing.time.time", return_value=start_time + 61): 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) self.login_user(user_1)
response = self.client_get(url, subdomain=test_subdomain) response = self.client_get(url, subdomain=test_subdomain)
self.assertEqual(response.status_code, 200) self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!") self.assertEqual(response.getvalue(), b"zulip!")
self.logout() self.logout()
# Confirm other cross-realm users can't read it. # Confirm other cross-realm users can't read it.
@ -784,7 +783,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
with self.assert_database_query_count(5): with self.assert_database_query_count(5):
response = self.client_get(url) response = self.client_get(url)
self.assertEqual(response.status_code, 200) self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!") self.assertEqual(response.getvalue(), b"zulip!")
self.logout() self.logout()
# Subscribed user who received the message should be able to view file # 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): with self.assert_database_query_count(6):
response = self.client_get(url) response = self.client_get(url)
self.assertEqual(response.status_code, 200) self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!") self.assertEqual(response.getvalue(), b"zulip!")
self.logout() self.logout()
def assert_cannot_access_file(user: UserProfile) -> None: def assert_cannot_access_file(user: UserProfile) -> None:
@ -845,7 +844,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
with self.assert_database_query_count(5): with self.assert_database_query_count(5):
response = self.client_get(url) response = self.client_get(url)
self.assertEqual(response.status_code, 200) self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!") self.assertEqual(response.getvalue(), b"zulip!")
self.logout() self.logout()
# Originally subscribed user should be able to view file # Originally subscribed user should be able to view file
@ -853,7 +852,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
with self.assert_database_query_count(6): with self.assert_database_query_count(6):
response = self.client_get(url) response = self.client_get(url)
self.assertEqual(response.status_code, 200) self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!") self.assertEqual(response.getvalue(), b"zulip!")
self.logout() self.logout()
# Subscribed user who did not receive the message should also be able to view file # 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): with self.assert_database_query_count(9):
response = self.client_get(url) response = self.client_get(url)
self.assertEqual(response.status_code, 200) self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!") self.assertEqual(response.getvalue(), b"zulip!")
self.logout() self.logout()
# It takes a few extra queries to verify access because of shared history. # 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): with self.assert_database_query_count(9):
response = self.client_get(url) response = self.client_get(url)
self.assertEqual(response.status_code, 200) 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): with self.assert_database_query_count(6):
self.assertTrue(validate_attachment_request(user, fp_path_id)) self.assertTrue(validate_attachment_request(user, fp_path_id))
@ -949,7 +948,7 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
for user in subscribed_users + unsubscribed_users: for user in subscribed_users + unsubscribed_users:
self.login_user(user) self.login_user(user)
response = self.client_get(url) response = self.client_get(url)
self.assert_streaming_content(response, b"zulip!") self.assertEqual(response.getvalue(), b"zulip!")
self.logout() self.logout()
def test_serve_local(self) -> None: def test_serve_local(self) -> None:
@ -1309,8 +1308,7 @@ class AvatarTest(UploadSerializeMixin, ZulipTestCase):
if rfname is not None: if rfname is not None:
response = self.client_get(url) response = self.client_get(url)
assert isinstance(response, StreamingHttpResponse) data = response.getvalue()
data = b"".join(response.streaming_content)
self.assertEqual(Image.open(io.BytesIO(data)).size, (100, 100)) self.assertEqual(Image.open(io.BytesIO(data)).size, (100, 100))
# Verify that the medium-size avatar was created # Verify that the medium-size avatar was created
@ -1590,8 +1588,7 @@ class RealmIconTest(UploadSerializeMixin, ZulipTestCase):
if rfname is not None: if rfname is not None:
response = self.client_get(url) response = self.client_get(url)
assert isinstance(response, StreamingHttpResponse) data = response.getvalue()
data = b"".join(response.streaming_content)
self.assertEqual(Image.open(io.BytesIO(data)).size, (100, 100)) self.assertEqual(Image.open(io.BytesIO(data)).size, (100, 100))
def test_invalid_icons(self) -> None: def test_invalid_icons(self) -> None:
@ -1773,8 +1770,7 @@ class RealmLogoTest(UploadSerializeMixin, ZulipTestCase):
if rfname is not None: if rfname is not None:
response = self.client_get(logo_url) response = self.client_get(logo_url)
assert isinstance(response, StreamingHttpResponse) data = response.getvalue()
data = b"".join(response.streaming_content)
# size should be 100 x 100 because thumbnail keeps aspect ratio # 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 # 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)) self.assertEqual(Image.open(io.BytesIO(data)).size, (100, 100))

View File

@ -5,7 +5,6 @@ from io import BytesIO, StringIO
from urllib.parse import urlparse from urllib.parse import urlparse
from django.conf import settings from django.conf import settings
from django.http.response import StreamingHttpResponse
from PIL import Image from PIL import Image
import zerver.lib.upload import zerver.lib.upload
@ -150,8 +149,7 @@ class LocalStorageTest(UploadSerializeMixin, ZulipTestCase):
# We get a resized avatar from it # We get a resized avatar from it
image_data = read_test_image_file("img.png") image_data = read_test_image_file("img.png")
resized_avatar = resize_avatar(image_data) resized_avatar = resize_avatar(image_data)
assert isinstance(result, StreamingHttpResponse) self.assertEqual(resized_avatar, result.getvalue())
self.assertEqual(resized_avatar, b"".join(result.streaming_content))
with self.settings(DEVELOPMENT=False): with self.settings(DEVELOPMENT=False):
# In production, this is an X-Accel-Redirect to the # In production, this is an X-Accel-Redirect to the