test_classes: Create a dedicate helper for query count check.

This adds a helper based on testing patterns of using the "queries_captured"
context manager with "assert_length" to check the number of queries
executed for preventing performance regression.

It explains the rationale of checking the query count through an
"AssertionError" and prints the queries captured as assert_length does,
but with a format optimized for displaying the queries in a more
readable manner.

Signed-off-by: Zixuan James Li <p359101898@gmail.com>
This commit is contained in:
Zixuan James Li 2022-10-15 16:47:40 -04:00 committed by Tim Abbott
parent 9117db8887
commit 46329a2710
21 changed files with 122 additions and 200 deletions

View File

@ -3,7 +3,6 @@ from unittest import mock
from django.utils.timezone import now as timezone_now
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import queries_captured
from zerver.models import Client, UserActivity, UserProfile, flush_per_request_caches
@ -33,23 +32,17 @@ class ActivityTest(ZulipTestCase):
user_profile.save(update_fields=["is_staff"])
flush_per_request_caches()
with queries_captured() as queries:
with self.assert_database_query_count(19):
result = self.client_get("/activity")
self.assertEqual(result.status_code, 200)
self.assert_length(queries, 19)
flush_per_request_caches()
with queries_captured() as queries:
with self.assert_database_query_count(8):
result = self.client_get("/realm_activity/zulip/")
self.assertEqual(result.status_code, 200)
self.assert_length(queries, 8)
iago = self.example_user("iago")
flush_per_request_caches()
with queries_captured() as queries:
with self.assert_database_query_count(5):
result = self.client_get(f"/user_activity/{iago.id}/")
self.assertEqual(result.status_code, 200)
self.assert_length(queries, 5)

View File

@ -435,8 +435,8 @@ requires a more subtle restructuring of the code.)
We try to prevent these bugs in our tests by using a context manager
called `queries_captured()` that captures the SQL queries used by
the backend during a particular operation. We make assertions about
those queries, often simply asserting that the number of queries is
below some threshold.
those queries, often simply by using the `assert_database_query_count`
that checks the number of queries.
### Event-based tests

View File

@ -70,7 +70,7 @@ from zerver.lib.test_console_output import (
tee_stderr_and_find_extra_console_output,
tee_stdout_and_find_extra_console_output,
)
from zerver.lib.test_helpers import find_key_by_email, instrument_url
from zerver.lib.test_helpers import find_key_by_email, instrument_url, queries_captured
from zerver.lib.topic import filter_by_topic_name_via_message
from zerver.lib.user_groups import get_system_user_group_for_user
from zerver.lib.users import get_api_key
@ -1123,6 +1123,33 @@ Output:
print(f"\nexpected length: {count}\nactual length: {actual_count}")
raise AssertionError(f"{str(type(items))} is of unexpected size!")
@contextmanager
def assert_database_query_count(
self, count: int, include_savepoints: bool = False, keep_cache_warm: bool = False
) -> Iterator[None]:
"""
This captures the queries executed and check the total number of queries.
Useful when minimizing unnecessary roundtrips to the database is important.
"""
with queries_captured(
include_savepoints=include_savepoints, keep_cache_warm=keep_cache_warm
) as queries:
yield
actual_count = len(queries)
if actual_count != count: # nocoverage
print("\nITEMS:\n")
for index, query in enumerate(queries):
print(f"#{index + 1}\nsql: {str(query['sql'])}\ntime: {query['time']}\n")
print(f"expected count: {count}\nactual count: {actual_count}")
raise AssertionError(
f"""
{count} queries expected, got {actual_count}.
This is a performance-critical code path, where we check
the number of database queries used in order to avoid accidental regressions.
If the new query is necessary, you should update this test,
and explain in the pull request why adding an additional query can't be avoided."""
)
def assert_json_error_contains(
self, result: "TestHttpResponse", msg_substring: str, status_code: int = 400
) -> None:

View File

@ -16,7 +16,7 @@ from zerver.lib.bot_config import ConfigError, get_bot_config
from zerver.lib.bot_lib import get_bot_handler
from zerver.lib.integrations import EMBEDDED_BOTS, WebhookIntegration
from zerver.lib.test_classes import UploadSerializeMixin, ZulipTestCase
from zerver.lib.test_helpers import avatar_disk_path, get_test_image_file, queries_captured
from zerver.lib.test_helpers import avatar_disk_path, get_test_image_file
from zerver.models import (
Realm,
Service,
@ -153,13 +153,11 @@ class BotTest(ZulipTestCase, UploadSerializeMixin):
self.assert_num_bots_equal(num_bots)
with queries_captured() as queries:
with self.assert_database_query_count(3):
users_result = self.client_get("/json/users")
self.assert_json_success(users_result)
self.assert_length(queries, 3)
def test_add_bot(self) -> None:
hamlet = self.example_user("hamlet")
self.login("hamlet")

View File

@ -21,7 +21,6 @@ from zerver.lib.cache import (
validate_cache_key,
)
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import queries_captured
from zerver.models import UserProfile, get_realm, get_system_bot, get_user, get_user_profile_by_id
@ -86,12 +85,11 @@ class CacheWithKeyDecoratorTest(ZulipTestCase):
hamlet = self.example_user("hamlet")
with patch("zerver.lib.cache.cache_set") as mock_set, self.assertLogs(level="WARNING") as m:
with queries_captured() as queries:
with self.assert_database_query_count(1):
result = get_user_function_with_bad_cache_keys(hamlet.id)
self.assert_length(m.output, 1)
self.assertEqual(result, hamlet)
self.assert_length(queries, 1)
mock_set.assert_not_called()
def test_cache_with_key_key_too_long(self) -> None:
@ -105,12 +103,11 @@ class CacheWithKeyDecoratorTest(ZulipTestCase):
hamlet = self.example_user("hamlet")
with patch("zerver.lib.cache.cache_set") as mock_set, self.assertLogs(level="WARNING") as m:
with queries_captured() as queries:
with self.assert_database_query_count(1):
result = get_user_function_with_bad_cache_keys(hamlet.id)
self.assert_length(m.output, 1)
self.assertEqual(result, hamlet)
self.assert_length(queries, 1)
mock_set.assert_not_called()
def test_cache_with_key_good_key(self) -> None:
@ -123,19 +120,17 @@ class CacheWithKeyDecoratorTest(ZulipTestCase):
hamlet = self.example_user("hamlet")
with queries_captured() as queries:
with self.assert_database_query_count(1):
result = get_user_function_with_good_cache_keys(hamlet.id)
self.assertEqual(result, hamlet)
self.assert_length(queries, 1)
# The previous function call should have cached the result correctly, so now
# no database queries should happen:
with queries_captured(keep_cache_warm=True) as queries_two:
with self.assert_database_query_count(0, keep_cache_warm=True):
result_two = get_user_function_with_good_cache_keys(hamlet.id)
self.assertEqual(result_two, hamlet)
self.assert_length(queries_two, 0)
def test_cache_with_key_none_values(self) -> None:
def cache_key_function(user_id: int) -> str:
@ -151,17 +146,15 @@ class CacheWithKeyDecoratorTest(ZulipTestCase):
last_user = UserProfile.objects.last()
assert last_user is not None
last_user_id = last_user.id
with queries_captured() as queries:
with self.assert_database_query_count(1):
result = get_user_function_can_return_none(last_user_id + 1)
self.assertEqual(result, None)
self.assert_length(queries, 1)
with queries_captured(keep_cache_warm=True) as queries:
with self.assert_database_query_count(0, keep_cache_warm=True):
result_two = get_user_function_can_return_none(last_user_id + 1)
self.assertEqual(result_two, None)
self.assert_length(queries, 0)
class SafeCacheFunctionsTest(ZulipTestCase):

View File

@ -12,7 +12,6 @@ from zerver.actions.custom_profile_fields import (
from zerver.lib.external_accounts import DEFAULT_EXTERNAL_ACCOUNTS
from zerver.lib.markdown import markdown_convert
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import queries_captured
from zerver.lib.types import ProfileDataElementUpdateDict, ProfileDataElementValue
from zerver.models import (
CustomProfileField,
@ -908,13 +907,11 @@ class ListCustomProfileFieldTest(CustomProfileFieldTestCase):
test_bot = self.create_test_bot("foo-bot", iago)
self.login_user(iago)
with queries_captured() as queries:
with self.assert_database_query_count(4):
response = self.client_get(
"/json/users", {"client_gravatar": "false", "include_custom_profile_fields": "true"}
)
self.assert_length(queries, 4)
raw_users_data = self.assert_json_success(response)["members"]
iago_raw_data = None

View File

@ -24,7 +24,7 @@ from zerver.lib.digest import (
from zerver.lib.message import get_last_message_id
from zerver.lib.streams import create_stream_if_needed
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import cache_tries_captured, queries_captured
from zerver.lib.test_helpers import cache_tries_captured
from zerver.models import (
Client,
Message,
@ -69,11 +69,9 @@ class TestDigestEmailMessages(ZulipTestCase):
# This code is run when we call `confirmation.models.create_confirmation_link`.
# To trigger this, we call the one_click_unsubscribe_link function below.
one_click_unsubscribe_link(othello, "digest")
with queries_captured() as queries:
with self.assert_database_query_count(9):
bulk_handle_digest_email([othello.id], cutoff)
self.assert_length(queries, 9)
self.assertEqual(mock_send_future_email.call_count, 1)
kwargs = mock_send_future_email.call_args[1]
self.assertEqual(kwargs["to_user_ids"], [othello.id])
@ -148,11 +146,9 @@ class TestDigestEmailMessages(ZulipTestCase):
# This code is run when we call `confirmation.models.create_confirmation_link`.
# To trigger this, we call the one_click_unsubscribe_link function below.
one_click_unsubscribe_link(polonius, "digest")
with queries_captured() as queries:
with self.assert_database_query_count(9):
bulk_handle_digest_email([polonius.id], cutoff)
self.assert_length(queries, 9)
self.assertEqual(mock_send_future_email.call_count, 1)
kwargs = mock_send_future_email.call_args[1]
self.assertEqual(kwargs["to_user_ids"], [polonius.id])
@ -213,11 +209,10 @@ class TestDigestEmailMessages(ZulipTestCase):
with mock.patch("zerver.lib.digest.send_future_email") as mock_send_future_email:
digest_user_ids = [user.id for user in digest_users]
with queries_captured() as queries:
with self.assert_database_query_count(12):
with cache_tries_captured() as cache_tries:
bulk_handle_digest_email(digest_user_ids, cutoff)
self.assert_length(queries, 12)
self.assert_length(cache_tries, 0)
self.assert_length(digest_users, mock_send_future_email.call_count)

View File

@ -18,12 +18,7 @@ from zerver.lib.events import fetch_initial_state_data
from zerver.lib.exceptions import AccessDeniedError
from zerver.lib.request import RequestVariableMissingError
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import (
HostRequestMock,
dummy_handler,
queries_captured,
stub_event_queue_user_events,
)
from zerver.lib.test_helpers import HostRequestMock, dummy_handler, stub_event_queue_user_events
from zerver.lib.users import get_api_key, get_raw_user_data
from zerver.models import (
Realm,
@ -1114,12 +1109,10 @@ class FetchQueriesTest(ZulipTestCase):
self.login_user(user)
flush_per_request_caches()
with queries_captured() as queries:
with self.assert_database_query_count(37):
with mock.patch("zerver.lib.events.always_want") as want_mock:
fetch_initial_state_data(user)
self.assert_length(queries, 37)
expected_counts = dict(
alert_words=1,
custom_profile_fields=1,
@ -1165,14 +1158,13 @@ class FetchQueriesTest(ZulipTestCase):
for event_type in sorted(wanted_event_types):
count = expected_counts[event_type]
flush_per_request_caches()
with queries_captured() as queries:
with self.assert_database_query_count(count):
if event_type == "update_message_flags":
event_types = ["update_message_flags", "message"]
else:
event_types = [event_type]
fetch_initial_state_data(user, event_types=event_types)
self.assert_length(queries, count)
class TestEventsRegisterAllPublicStreamsDefaults(ZulipTestCase):

View File

@ -9,7 +9,7 @@ from zerver.actions.users import do_change_can_create_users, do_change_user_role
from zerver.lib.exceptions import JsonableError
from zerver.lib.streams import access_stream_for_send_message
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import most_recent_message, queries_captured
from zerver.lib.test_helpers import most_recent_message
from zerver.lib.users import is_administrator_role
from zerver.models import (
UserProfile,
@ -357,8 +357,8 @@ class TestQueryCounts(ZulipTestCase):
def test_capturing_queries(self) -> None:
# It's a common pitfall in Django to accidentally perform
# database queries in a loop, due to lazy evaluation of
# foreign keys. We use the queries_captured context manager to
# ensure our query count is predictable.
# foreign keys. We use the assert_database_query_count
# context manager to ensure our query count is predictable.
#
# When a test containing one of these query count assertions
# fails, we'll want to understand the new queries and whether
@ -368,16 +368,13 @@ class TestQueryCounts(ZulipTestCase):
hamlet = self.example_user("hamlet")
cordelia = self.example_user("cordelia")
with queries_captured() as queries:
with self.assert_database_query_count(15):
self.send_personal_message(
from_user=hamlet,
to_user=cordelia,
content="hello there!",
)
# The assert_length helper is another useful extra from ZulipTestCase.
self.assert_length(queries, 15)
class TestDevelopmentEmailsLog(ZulipTestCase):
# We have development specific utilities that automate common tasks

View File

@ -245,7 +245,7 @@ class HomeTest(ZulipTestCase):
# Verify succeeds once logged-in
flush_per_request_caches()
with queries_captured() as queries:
with self.assert_database_query_count(47):
with patch("zerver.lib.cache.cache_set") as cache_mock:
result = self._get_home_page(stream="Denmark")
self.check_rendered_logged_in_app(result)
@ -253,7 +253,6 @@ class HomeTest(ZulipTestCase):
set(result["Cache-Control"].split(", ")), {"must-revalidate", "no-store", "no-cache"}
)
self.assert_length(queries, 47)
self.assert_length(cache_mock.call_args_list, 5)
html = result.content.decode()
@ -392,12 +391,11 @@ class HomeTest(ZulipTestCase):
# Verify number of queries for Realm admin isn't much higher than for normal users.
self.login("iago")
flush_per_request_caches()
with queries_captured() as queries:
with self.assert_database_query_count(44):
with patch("zerver.lib.cache.cache_set") as cache_mock:
result = self._get_home_page()
self.check_rendered_logged_in_app(result)
self.assert_length(cache_mock.call_args_list, 6)
self.assert_length(queries, 44)
def test_num_queries_with_streams(self) -> None:
main_user = self.example_user("hamlet")
@ -425,11 +423,9 @@ class HomeTest(ZulipTestCase):
# Then for the second page load, measure the number of queries.
flush_per_request_caches()
with queries_captured() as queries2:
with self.assert_database_query_count(42):
result = self._get_home_page()
self.assert_length(queries2, 42)
# Do a sanity check that our new streams were in the payload.
html = result.content.decode()
self.assertIn("test_stream_7", html)

View File

@ -7,7 +7,7 @@ from zerver.lib.cache import cache_delete, to_dict_cache_key_id
from zerver.lib.markdown import version as markdown_version
from zerver.lib.message import MessageDict, messages_for_ids, sew_messages_and_reactions
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import make_client, queries_captured
from zerver.lib.test_helpers import make_client
from zerver.lib.topic import TOPIC_LINKS
from zerver.lib.types import DisplayRecipientT, UserDisplayRecipient
from zerver.models import (
@ -177,13 +177,12 @@ class MessageDictTest(ZulipTestCase):
self.assertTrue(num_ids >= 600)
flush_per_request_caches()
with queries_captured() as queries:
with self.assert_database_query_count(7):
rows = list(MessageDict.get_raw_db_rows(ids))
objs = [MessageDict.build_dict_from_raw_db_row(row) for row in rows]
MessageDict.post_process_dicts(objs, apply_markdown=False, client_gravatar=False)
self.assert_length(queries, 7)
self.assert_length(rows, num_ids)
def test_applying_markdown(self) -> None:

View File

@ -273,18 +273,16 @@ class EditMessageTest(EditMessageTestCase):
]
# Check number of queries performed
with queries_captured() as queries:
MessageDict.to_dict_uncached(messages)
# 1 query for realm_id per message = 3
# 1 query each for reactions & submessage for all messages = 2
self.assert_length(queries, 5)
with self.assert_database_query_count(5):
MessageDict.to_dict_uncached(messages)
realm_id = 2 # Fetched from stream object
# Check number of queries performed with realm_id
with queries_captured() as queries:
MessageDict.to_dict_uncached(messages, realm_id)
# 1 query each for reactions & submessage for all messages = 2
self.assert_length(queries, 2)
with self.assert_database_query_count(2):
MessageDict.to_dict_uncached(messages, realm_id)
def test_save_message(self) -> None:
"""This is also tested by a client test, but here we can verify
@ -1294,7 +1292,9 @@ class EditMessageTest(EditMessageTestCase):
users_to_be_notified = list(map(notify, [hamlet.id, cordelia.id, aaron.id]))
change_all_topic_name = "Topic 1 edited"
with queries_captured() as queries:
# This code path adds 9 (1 + 4/user with muted topics) + 1 to
# the number of database queries for moving a topic.
with self.assert_database_query_count(19):
check_update_message(
user_profile=hamlet,
message_id=message_id,
@ -1305,9 +1305,6 @@ class EditMessageTest(EditMessageTestCase):
send_notification_to_new_thread=False,
content=None,
)
# This code path adds 9 (1 + 4/user with muted topics) + 1 to
# the number of database queries for moving a topic.
self.assert_length(queries, 19)
for muting_user in get_users_muting_topic(stream.id, change_all_topic_name):
for user in users_to_be_notified:
@ -1381,7 +1378,7 @@ class EditMessageTest(EditMessageTestCase):
set_topic_mutes(desdemona, muted_topics)
set_topic_mutes(cordelia, muted_topics)
with queries_captured() as queries:
with self.assert_database_query_count(31):
check_update_message(
user_profile=desdemona,
message_id=message_id,
@ -1391,7 +1388,6 @@ class EditMessageTest(EditMessageTestCase):
send_notification_to_new_thread=False,
content=None,
)
self.assert_length(queries, 31)
self.assertFalse(topic_is_muted(desdemona, stream.id, "New topic"))
self.assertFalse(topic_is_muted(cordelia, stream.id, "New topic"))
@ -1413,7 +1409,7 @@ class EditMessageTest(EditMessageTestCase):
set_topic_mutes(desdemona, muted_topics)
set_topic_mutes(cordelia, muted_topics)
with queries_captured() as queries:
with self.assert_database_query_count(33):
check_update_message(
user_profile=desdemona,
message_id=message_id,
@ -1423,7 +1419,6 @@ class EditMessageTest(EditMessageTestCase):
send_notification_to_new_thread=False,
content=None,
)
self.assert_length(queries, 33)
# Cordelia is not subscribed to the private stream, so
# Cordelia should have had the topic unmuted, while Desdemona
@ -1447,7 +1442,7 @@ class EditMessageTest(EditMessageTestCase):
set_topic_mutes(desdemona, muted_topics)
set_topic_mutes(cordelia, muted_topics)
with queries_captured() as queries:
with self.assert_database_query_count(31):
check_update_message(
user_profile=desdemona,
message_id=message_id,
@ -1458,7 +1453,6 @@ class EditMessageTest(EditMessageTestCase):
send_notification_to_new_thread=False,
content=None,
)
self.assert_length(queries, 31)
self.assertFalse(topic_is_muted(desdemona, stream.id, "New topic 2"))
self.assertFalse(topic_is_muted(cordelia, stream.id, "New topic 2"))
@ -1471,7 +1465,7 @@ class EditMessageTest(EditMessageTestCase):
second_message_id = self.send_stream_message(
hamlet, stream_name, topic_name="changed topic name", content="Second message"
)
with queries_captured() as queries:
with self.assert_database_query_count(25):
check_update_message(
user_profile=desdemona,
message_id=second_message_id,
@ -1482,7 +1476,6 @@ class EditMessageTest(EditMessageTestCase):
send_notification_to_new_thread=False,
content=None,
)
self.assert_length(queries, 25)
self.assertTrue(topic_is_muted(desdemona, new_public_stream.id, "changed topic name"))
self.assertTrue(topic_is_muted(cordelia, new_public_stream.id, "changed topic name"))
@ -2236,7 +2229,7 @@ class EditMessageTest(EditMessageTestCase):
"iago", "test move stream", "new stream", "test"
)
with queries_captured() as queries, cache_tries_captured() as cache_tries:
with self.assert_database_query_count(53), cache_tries_captured() as cache_tries:
result = self.client_patch(
f"/json/messages/{msg_id}",
{
@ -2246,7 +2239,6 @@ class EditMessageTest(EditMessageTestCase):
"topic": "new topic",
},
)
self.assert_length(queries, 53)
self.assert_length(cache_tries, 13)
messages = get_topic_messages(user_profile, old_stream, "test")

View File

@ -21,7 +21,7 @@ from zerver.lib.message import (
get_raw_unread_data,
)
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import get_subscription, queries_captured
from zerver.lib.test_helpers import get_subscription
from zerver.lib.user_topics import add_topic_mute
from zerver.models import (
Message,
@ -1332,9 +1332,8 @@ class MessageAccessTests(ZulipTestCase):
Message.objects.select_related().get(id=message_id) for message_id in message_ids
]
with queries_captured() as queries:
with self.assert_database_query_count(2):
filtered_messages = bulk_access_messages(later_subscribed_user, messages, stream=stream)
self.assert_length(queries, 2)
# Message sent before subscribing wouldn't be accessible by later
# subscribed user as stream has protected history
@ -1349,9 +1348,8 @@ class MessageAccessTests(ZulipTestCase):
acting_user=self.example_user("cordelia"),
)
with queries_captured() as queries:
with self.assert_database_query_count(2):
filtered_messages = bulk_access_messages(later_subscribed_user, messages, stream=stream)
self.assert_length(queries, 2)
# Message sent before subscribing are accessible by 8user as stream
# don't have protected history
@ -1360,9 +1358,8 @@ class MessageAccessTests(ZulipTestCase):
# Testing messages accessibility for an unsubscribed user
unsubscribed_user = self.example_user("ZOE")
with queries_captured() as queries:
with self.assert_database_query_count(2):
filtered_messages = bulk_access_messages(unsubscribed_user, messages, stream=stream)
self.assert_length(queries, 2)
self.assert_length(filtered_messages, 0)
@ -1394,17 +1391,15 @@ class MessageAccessTests(ZulipTestCase):
]
# All public stream messages are always accessible
with queries_captured() as queries:
with self.assert_database_query_count(2):
filtered_messages = bulk_access_messages(later_subscribed_user, messages, stream=stream)
self.assert_length(filtered_messages, 2)
self.assert_length(queries, 2)
unsubscribed_user = self.example_user("ZOE")
with queries_captured() as queries:
with self.assert_database_query_count(2):
filtered_messages = bulk_access_messages(unsubscribed_user, messages, stream=stream)
self.assert_length(filtered_messages, 2)
self.assert_length(queries, 2)
class PersonalMessagesFlagTest(ZulipTestCase):

View File

@ -42,7 +42,6 @@ from zerver.lib.test_helpers import (
message_stream_count,
most_recent_message,
most_recent_usermessage,
queries_captured,
reset_emails_in_zulip_realm,
)
from zerver.lib.timestamp import convert_to_UTC, datetime_to_timestamp
@ -1622,7 +1621,7 @@ class StreamMessagesTest(ZulipTestCase):
# during the course of the unit test.
flush_per_request_caches()
cache_delete(get_stream_cache_key(stream_name, realm.id))
with queries_captured() as queries:
with self.assert_database_query_count(13):
check_send_stream_message(
sender=sender,
client=sending_client,
@ -1631,8 +1630,6 @@ class StreamMessagesTest(ZulipTestCase):
body=content,
)
self.assert_length(queries, 13)
def test_stream_message_dict(self) -> None:
user_profile = self.example_user("iago")
self.subscribe(user_profile, "Denmark")

View File

@ -19,7 +19,7 @@ from zerver.lib.retention import (
restore_retention_policy_deletions_for_stream,
)
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import queries_captured, zulip_reaction_info
from zerver.lib.test_helpers import zulip_reaction_info
from zerver.lib.upload import create_attachment
from zerver.models import (
ArchivedAttachment,
@ -1053,10 +1053,9 @@ class TestDoDeleteMessages(ZulipTestCase):
message_ids = [self.send_stream_message(cordelia, "Verona", str(i)) for i in range(0, 10)]
messages = Message.objects.filter(id__in=message_ids)
with queries_captured() as queries:
with self.assert_database_query_count(19):
do_delete_messages(realm, messages)
self.assertFalse(Message.objects.filter(id__in=message_ids).exists())
self.assert_length(queries, 19)
archived_messages = ArchivedMessage.objects.filter(id__in=message_ids)
self.assertEqual(archived_messages.count(), len(message_ids))

View File

@ -956,11 +956,10 @@ class LoginTest(ZulipTestCase):
flush_per_request_caches()
ContentType.objects.clear_cache()
with queries_captured() as queries, cache_tries_captured() as cache_tries:
# Ensure the number of queries we make is not O(streams)
with self.assert_database_query_count(96), cache_tries_captured() as cache_tries:
with self.captureOnCommitCallbacks(execute=True):
self.register(self.nonreg_email("test"), "test")
# Ensure the number of queries we make is not O(streams)
self.assert_length(queries, 96)
# We can probably avoid a couple cache hits here, but there doesn't
# seem to be any O(N) behavior. Some of the cache hits are related

View File

@ -17,12 +17,7 @@ from zerver.lib.soft_deactivation import (
)
from zerver.lib.stream_subscription import get_subscriptions_for_send_message
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import (
get_subscription,
get_user_messages,
make_client,
queries_captured,
)
from zerver.lib.test_helpers import get_subscription, get_user_messages, make_client
from zerver.models import (
AlertWord,
Client,
@ -340,9 +335,8 @@ class SoftDeactivationMessageTest(ZulipTestCase):
idle_user_msg_list = get_user_messages(long_term_idle_user)
idle_user_msg_count = len(idle_user_msg_list)
self.assertNotEqual(idle_user_msg_list[-1].content, message)
with queries_captured() as queries:
with self.assert_database_query_count(8):
reactivate_user_if_soft_deactivated(long_term_idle_user)
self.assert_length(queries, 8)
self.assertFalse(long_term_idle_user.long_term_idle)
self.assertEqual(
last_realm_audit_log_entry(RealmAuditLog.USER_SOFT_ACTIVATED).modified_user,
@ -402,9 +396,8 @@ class SoftDeactivationMessageTest(ZulipTestCase):
idle_user_msg_list = get_user_messages(long_term_idle_user)
idle_user_msg_count = len(idle_user_msg_list)
self.assertNotEqual(idle_user_msg_list[-1], sent_message)
with queries_captured() as queries:
with self.assert_database_query_count(6):
add_missing_messages(long_term_idle_user)
self.assert_length(queries, 6)
idle_user_msg_list = get_user_messages(long_term_idle_user)
self.assert_length(idle_user_msg_list, idle_user_msg_count + 1)
self.assertEqual(idle_user_msg_list[-1], sent_message)
@ -419,9 +412,8 @@ class SoftDeactivationMessageTest(ZulipTestCase):
idle_user_msg_list = get_user_messages(long_term_idle_user)
idle_user_msg_count = len(idle_user_msg_list)
self.assertNotEqual(idle_user_msg_list[-1], sent_message)
with queries_captured() as queries:
with self.assert_database_query_count(7):
add_missing_messages(long_term_idle_user)
self.assert_length(queries, 7)
idle_user_msg_list = get_user_messages(long_term_idle_user)
self.assert_length(idle_user_msg_list, idle_user_msg_count + 1)
self.assertEqual(idle_user_msg_list[-1], sent_message)
@ -444,9 +436,8 @@ class SoftDeactivationMessageTest(ZulipTestCase):
idle_user_msg_count = len(idle_user_msg_list)
for sent_message in sent_message_list:
self.assertNotEqual(idle_user_msg_list.pop(), sent_message)
with queries_captured() as queries:
with self.assert_database_query_count(6):
add_missing_messages(long_term_idle_user)
self.assert_length(queries, 6)
idle_user_msg_list = get_user_messages(long_term_idle_user)
self.assert_length(idle_user_msg_list, idle_user_msg_count + 2)
for sent_message in sent_message_list:
@ -476,9 +467,8 @@ class SoftDeactivationMessageTest(ZulipTestCase):
idle_user_msg_count = len(idle_user_msg_list)
for sent_message in sent_message_list:
self.assertNotEqual(idle_user_msg_list.pop(), sent_message)
with queries_captured() as queries:
with self.assert_database_query_count(6):
add_missing_messages(long_term_idle_user)
self.assert_length(queries, 6)
idle_user_msg_list = get_user_messages(long_term_idle_user)
self.assert_length(idle_user_msg_list, idle_user_msg_count + 2)
for sent_message in sent_message_list:
@ -509,11 +499,10 @@ class SoftDeactivationMessageTest(ZulipTestCase):
idle_user_msg_list = get_user_messages(long_term_idle_user)
idle_user_msg_count = len(idle_user_msg_list)
self.assertEqual(idle_user_msg_list[-1].id, sent_message_id)
with queries_captured() as queries:
add_missing_messages(long_term_idle_user)
# There are no streams to fetch missing messages from, so
# the Message.objects query will be avoided.
self.assert_length(queries, 4)
with self.assert_database_query_count(4):
add_missing_messages(long_term_idle_user)
idle_user_msg_list = get_user_messages(long_term_idle_user)
# No new UserMessage rows should have been created.
self.assert_length(idle_user_msg_list, idle_user_msg_count)
@ -538,9 +527,8 @@ class SoftDeactivationMessageTest(ZulipTestCase):
idle_user_msg_count = len(idle_user_msg_list)
for sent_message in sent_message_list:
self.assertNotEqual(idle_user_msg_list.pop(), sent_message)
with queries_captured() as queries:
with self.assert_database_query_count(6):
add_missing_messages(long_term_idle_user)
self.assert_length(queries, 6)
idle_user_msg_list = get_user_messages(long_term_idle_user)
self.assert_length(idle_user_msg_list, idle_user_msg_count + 2)
for sent_message in sent_message_list:
@ -576,9 +564,8 @@ class SoftDeactivationMessageTest(ZulipTestCase):
idle_user_msg_list = get_user_messages(long_term_idle_user)
idle_user_msg_count = len(idle_user_msg_list)
with queries_captured() as queries:
with self.assert_database_query_count(10):
add_missing_messages(long_term_idle_user)
self.assert_length(queries, 10)
idle_user_msg_list = get_user_messages(long_term_idle_user)
self.assert_length(idle_user_msg_list, idle_user_msg_count + num_new_messages)
long_term_idle_user.refresh_from_db()

View File

@ -74,7 +74,6 @@ from zerver.lib.test_helpers import (
get_subscription,
most_recent_message,
most_recent_usermessage,
queries_captured,
reset_emails_in_zulip_realm,
)
from zerver.lib.types import (
@ -2216,7 +2215,7 @@ class StreamAdminTest(ZulipTestCase):
for user in other_sub_users:
self.subscribe(user, stream_name)
with queries_captured() as queries:
with self.assert_database_query_count(query_count):
with cache_tries_captured() as cache_tries:
result = self.client_delete(
"/json/users/me/subscriptions",
@ -2225,7 +2224,6 @@ class StreamAdminTest(ZulipTestCase):
"principals": orjson.dumps(principals).decode(),
},
)
self.assert_length(queries, query_count)
if cache_count is not None:
self.assert_length(cache_tries, cache_count)
@ -4225,13 +4223,12 @@ class SubscriptionAPITest(ZulipTestCase):
events: List[Mapping[str, Any]] = []
flush_per_request_caches()
with self.tornado_redirected_to_list(events, expected_num_events=5):
with queries_captured() as queries:
with self.assert_database_query_count(37):
self.common_subscribe_to_streams(
self.test_user,
streams_to_sub,
dict(principals=orjson.dumps([user1.id, user2.id]).decode()),
)
self.assert_length(queries, 37)
for ev in [x for x in events if x["event"]["type"] not in ("message", "stream")]:
if ev["event"]["op"] == "add":
@ -4250,13 +4247,12 @@ class SubscriptionAPITest(ZulipTestCase):
# Now add ourselves
with self.tornado_redirected_to_list(events, expected_num_events=2):
with queries_captured() as queries:
with self.assert_database_query_count(12):
self.common_subscribe_to_streams(
self.test_user,
streams_to_sub,
dict(principals=orjson.dumps([self.test_user.id]).decode()),
)
self.assert_length(queries, 12)
add_event, add_peer_event = events
self.assertEqual(add_event["event"]["type"], "subscription")
@ -4518,7 +4514,7 @@ class SubscriptionAPITest(ZulipTestCase):
# Sends 3 peer-remove events and 2 unsubscribe events.
events: List[Mapping[str, Any]] = []
with self.tornado_redirected_to_list(events, expected_num_events=5):
with queries_captured() as query_count:
with self.assert_database_query_count(16):
with cache_tries_captured() as cache_count:
bulk_remove_subscriptions(
realm,
@ -4527,7 +4523,6 @@ class SubscriptionAPITest(ZulipTestCase):
acting_user=None,
)
self.assert_length(query_count, 16)
self.assert_length(cache_count, 3)
peer_events = [e for e in events if e["event"].get("op") == "peer_remove"]
@ -4577,7 +4572,7 @@ class SubscriptionAPITest(ZulipTestCase):
# any tornado subscription events
events: List[Mapping[str, Any]] = []
with self.tornado_redirected_to_list(events, expected_num_events=0):
with queries_captured() as queries:
with self.assert_database_query_count(4):
self.common_subscribe_to_streams(
mit_user,
stream_names,
@ -4585,7 +4580,6 @@ class SubscriptionAPITest(ZulipTestCase):
subdomain="zephyr",
allow_fail=True,
)
self.assert_length(queries, 4)
with self.tornado_redirected_to_list(events, expected_num_events=0):
bulk_remove_subscriptions(
@ -4627,7 +4621,7 @@ class SubscriptionAPITest(ZulipTestCase):
test_user_ids = [user.id for user in test_users]
with queries_captured() as queries:
with self.assert_database_query_count(19):
with cache_tries_captured() as cache_tries:
with mock.patch("zerver.views.streams.send_messages_for_new_subscribers"):
self.common_subscribe_to_streams(
@ -4638,7 +4632,6 @@ class SubscriptionAPITest(ZulipTestCase):
# The only known O(N) behavior here is that we call
# principal_to_user_profile for each of our users.
self.assert_length(queries, 19)
self.assert_length(cache_tries, 4)
def test_subscriptions_add_for_principal(self) -> None:
@ -5102,29 +5095,27 @@ class SubscriptionAPITest(ZulipTestCase):
]
# Test creating a public stream when realm does not have a notification stream.
with queries_captured() as queries:
with self.assert_database_query_count(37):
self.common_subscribe_to_streams(
self.test_user,
[new_streams[0]],
dict(principals=orjson.dumps([user1.id, user2.id]).decode()),
)
self.assert_length(queries, 37)
# Test creating private stream.
with queries_captured() as queries:
with self.assert_database_query_count(36):
self.common_subscribe_to_streams(
self.test_user,
[new_streams[1]],
dict(principals=orjson.dumps([user1.id, user2.id]).decode()),
invite_only=True,
)
self.assert_length(queries, 36)
# Test creating a public stream with announce when realm has a notification stream.
notifications_stream = get_stream(self.streams[0], self.test_realm)
self.test_realm.notifications_stream_id = notifications_stream.id
self.test_realm.save()
with queries_captured() as queries:
with self.assert_database_query_count(45):
self.common_subscribe_to_streams(
self.test_user,
[new_streams[2]],
@ -5133,7 +5124,6 @@ class SubscriptionAPITest(ZulipTestCase):
principals=orjson.dumps([user1.id, user2.id]).decode(),
),
)
self.assert_length(queries, 45)
class GetStreamsTest(ZulipTestCase):
@ -5544,13 +5534,12 @@ class GetSubscribersTest(ZulipTestCase):
polonius.id,
]
with queries_captured() as queries:
with self.assert_database_query_count(46):
self.common_subscribe_to_streams(
self.user_profile,
streams,
dict(principals=orjson.dumps(users_to_subscribe).decode()),
)
self.assert_length(queries, 46)
msg = f"""
@**King Hamlet|{hamlet.id}** subscribed you to the following streams:
@ -5593,7 +5582,7 @@ class GetSubscribersTest(ZulipTestCase):
for user in [cordelia, othello, polonius]:
self.assert_user_got_subscription_notification(user, msg)
with queries_captured() as queries:
with self.assert_database_query_count(4):
subscribed_streams, _ = gather_subscriptions(
self.user_profile, include_subscribers=True
)
@ -5602,7 +5591,6 @@ class GetSubscribersTest(ZulipTestCase):
if not sub["name"].startswith("stream_"):
continue
self.assert_length(sub["subscribers"], len(users_to_subscribe))
self.assert_length(queries, 4)
def test_never_subscribed_streams(self) -> None:
"""
@ -5669,11 +5657,10 @@ class GetSubscribersTest(ZulipTestCase):
create_private_streams()
def get_never_subscribed() -> List[NeverSubscribedStreamDict]:
with queries_captured() as queries:
with self.assert_database_query_count(4):
sub_data = gather_subscriptions_helper(self.user_profile)
self.verify_sub_fields(sub_data)
never_subscribed = sub_data.never_subscribed
self.assert_length(queries, 4)
# Ignore old streams.
never_subscribed = [dct for dct in never_subscribed if dct["name"].startswith("test_")]
@ -5834,7 +5821,7 @@ class GetSubscribersTest(ZulipTestCase):
subdomain="zephyr",
)
with queries_captured() as queries:
with self.assert_database_query_count(4):
subscribed_streams, _ = gather_subscriptions(mit_user_profile, include_subscribers=True)
self.assertGreaterEqual(len(subscribed_streams), 2)
@ -5845,7 +5832,6 @@ class GetSubscribersTest(ZulipTestCase):
self.assert_length(sub["subscribers"], len(users_to_subscribe))
else:
self.assert_length(sub["subscribers"], 0)
self.assert_length(queries, 4)
def test_nonsubscriber(self) -> None:
"""

View File

@ -3,7 +3,6 @@ from typing import Any, List, Mapping
import orjson
from zerver.lib.test_classes import ZulipTestCase
from zerver.lib.test_helpers import queries_captured
from zerver.models import Huddle, get_huddle_hash
@ -148,13 +147,12 @@ class TypingHappyPathTestPMs(ZulipTestCase):
)
events: List[Mapping[str, Any]] = []
with queries_captured() as queries:
with self.assert_database_query_count(4):
with self.tornado_redirected_to_list(events, expected_num_events=1):
result = self.api_post(sender, "/api/v1/typing", params)
self.assert_json_success(result)
self.assert_length(events, 1)
self.assert_length(queries, 4)
event = events[0]["event"]
event_recipient_emails = {user["email"] for user in event["recipients"]}
@ -185,12 +183,11 @@ class TypingHappyPathTestPMs(ZulipTestCase):
op="start",
)
with queries_captured() as queries:
with self.assert_database_query_count(5):
with self.tornado_redirected_to_list(events, expected_num_events=1):
result = self.api_post(sender, "/api/v1/typing", params)
self.assert_json_success(result)
self.assert_length(events, 1)
self.assert_length(queries, 5)
# We should not be adding new Huddles just because
# a user started typing in the compose box. Let's
@ -366,12 +363,11 @@ class TypingHappyPathTestStreams(ZulipTestCase):
)
events: List[Mapping[str, Any]] = []
with queries_captured() as queries:
with self.assert_database_query_count(5):
with self.tornado_redirected_to_list(events, expected_num_events=1):
result = self.api_post(sender, "/api/v1/typing", params)
self.assert_json_success(result)
self.assert_length(events, 1)
self.assert_length(queries, 5)
event = events[0]["event"]
event_user_ids = set(events[0]["users"])
@ -402,12 +398,11 @@ class TypingHappyPathTestStreams(ZulipTestCase):
)
events: List[Mapping[str, Any]] = []
with queries_captured() as queries:
with self.assert_database_query_count(5):
with self.tornado_redirected_to_list(events, expected_num_events=1):
result = self.api_post(sender, "/api/v1/typing", params)
self.assert_json_success(result)
self.assert_length(events, 1)
self.assert_length(queries, 5)
event = events[0]["event"]
event_user_ids = set(events[0]["users"])

View File

@ -42,7 +42,6 @@ from zerver.lib.test_helpers import (
avatar_disk_path,
create_s3_buckets,
get_test_image_file,
queries_captured,
read_test_image_file,
use_s3_backend,
)
@ -757,21 +756,19 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
# Owner user should be able to view file
self.login_user(hamlet)
with queries_captured() as queries:
with self.assert_database_query_count(5):
response = self.client_get(uri)
self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!")
self.logout()
self.assert_length(queries, 5)
# Subscribed user who received the message should be able to view file
self.login_user(cordelia)
with queries_captured() as queries:
with self.assert_database_query_count(6):
response = self.client_get(uri)
self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!")
self.logout()
self.assert_length(queries, 6)
def assert_cannot_access_file(user: UserProfile) -> None:
response = self.api_get(user, uri)
@ -820,39 +817,35 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
# Owner user should be able to view file
self.login_user(user)
with queries_captured() as queries:
with self.assert_database_query_count(5):
response = self.client_get(uri)
self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!")
self.logout()
self.assert_length(queries, 5)
# Originally subscribed user should be able to view file
self.login_user(polonius)
with queries_captured() as queries:
with self.assert_database_query_count(6):
response = self.client_get(uri)
self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!")
self.logout()
self.assert_length(queries, 6)
# Subscribed user who did not receive the message should also be able to view file
self.login_user(late_subscribed_user)
with queries_captured() as queries:
with self.assert_database_query_count(9):
response = self.client_get(uri)
self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!")
self.logout()
# It takes a few extra queries to verify access because of shared history.
self.assert_length(queries, 9)
def assert_cannot_access_file(user: UserProfile) -> None:
self.login_user(user)
with queries_captured() as queries:
# It takes a few extra queries to verify lack of access with shared history.
with self.assert_database_query_count(8):
response = self.client_get(uri)
self.assertEqual(response.status_code, 403)
# It takes a few extra queries to verify lack of access with shared history.
self.assert_length(queries, 8)
self.assert_in_response("You are not authorized to view this file.", response)
self.logout()
@ -891,25 +884,22 @@ class FileUploadTest(UploadSerializeMixin, ZulipTestCase):
user = self.example_user("aaron")
self.login_user(user)
with queries_captured() as queries:
with self.assert_database_query_count(8):
response = self.client_get(uri)
self.assertEqual(response.status_code, 403)
self.assert_in_response("You are not authorized to view this file.", response)
self.assert_length(queries, 8)
self.subscribe(user, "test-subscribe 1")
self.subscribe(user, "test-subscribe 2")
with queries_captured() as queries:
# If we were accidentally one query per message, this would be 20+
with self.assert_database_query_count(9):
response = self.client_get(uri)
self.assertEqual(response.status_code, 200)
self.assert_streaming_content(response, b"zulip!")
# If we were accidentally one query per message, this would be 20+
self.assert_length(queries, 9)
with queries_captured() as queries:
with self.assert_database_query_count(6):
self.assertTrue(validate_attachment_request(user, fp_path_id))
self.assert_length(queries, 6)
self.logout()

View File

@ -41,7 +41,6 @@ from zerver.lib.test_helpers import (
cache_tries_captured,
get_subscription,
get_test_image_file,
queries_captured,
reset_emails_in_zulip_realm,
simulated_empty_cache,
)
@ -805,7 +804,7 @@ class QueryCountTest(ZulipTestCase):
events: List[Mapping[str, Any]] = []
with queries_captured() as queries:
with self.assert_database_query_count(91):
with cache_tries_captured() as cache_tries:
with self.tornado_redirected_to_list(events, expected_num_events=11):
fred = do_create_user(
@ -817,8 +816,6 @@ class QueryCountTest(ZulipTestCase):
acting_user=None,
)
self.assert_length(queries, 91)
self.assert_length(cache_tries, 28)
peer_add_events = [event for event in events if event["event"].get("op") == "peer_add"]
@ -1320,12 +1317,11 @@ class UserProfileTest(ZulipTestCase):
# Subscribe to the stream.
self.subscribe(iago, stream.name)
with queries_captured() as queries:
with self.assert_database_query_count(6):
result = orjson.loads(
self.client_get(f"/json/users/{iago.id}/subscriptions/{stream.id}").content
)
self.assert_length(queries, 6)
self.assertTrue(result["is_subscribed"])
# Logging in with a Guest user.
@ -2031,11 +2027,10 @@ class GetProfileTest(ZulipTestCase):
"""
realm = get_realm("zulip")
email = self.example_user("hamlet").email
with queries_captured() as queries:
with self.assert_database_query_count(1):
with simulated_empty_cache() as cache_queries:
user_profile = get_user(email, realm)
self.assert_length(queries, 1)
self.assert_length(cache_queries, 1)
self.assertEqual(user_profile.email, email)