From 46329a2710c0e8d78c3e516f18050361e5d2b5da Mon Sep 17 00:00:00 2001
From: Zixuan James Li <p359101898@gmail.com>
Date: Sat, 15 Oct 2022 16:47:40 -0400
Subject: [PATCH] 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>
---
 analytics/tests/test_activity_views.py   | 13 ++------
 docs/testing/testing-with-django.md      |  4 +--
 zerver/lib/test_classes.py               | 29 ++++++++++++++++-
 zerver/tests/test_bots.py                |  6 ++--
 zerver/tests/test_cache.py               | 19 ++++-------
 zerver/tests/test_custom_profile_data.py |  5 +--
 zerver/tests/test_digest.py              | 13 +++-----
 zerver/tests/test_event_system.py        | 14 ++-------
 zerver/tests/test_example.py             | 11 +++----
 zerver/tests/test_home.py                | 10 ++----
 zerver/tests/test_message_dict.py        |  5 ++-
 zerver/tests/test_message_edit.py        | 32 +++++++------------
 zerver/tests/test_message_flags.py       | 17 ++++------
 zerver/tests/test_message_send.py        |  5 +--
 zerver/tests/test_retention.py           |  5 ++-
 zerver/tests/test_signup.py              |  5 ++-
 zerver/tests/test_soft_deactivation.py   | 33 ++++++-------------
 zerver/tests/test_subs.py                | 40 ++++++++----------------
 zerver/tests/test_typing.py              | 13 +++-----
 zerver/tests/test_upload.py              | 32 +++++++------------
 zerver/tests/test_users.py               | 11 ++-----
 21 files changed, 122 insertions(+), 200 deletions(-)

diff --git a/analytics/tests/test_activity_views.py b/analytics/tests/test_activity_views.py
index 63c6ce8196..66cc5a3757 100644
--- a/analytics/tests/test_activity_views.py
+++ b/analytics/tests/test_activity_views.py
@@ -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)
diff --git a/docs/testing/testing-with-django.md b/docs/testing/testing-with-django.md
index d62efa20dc..1a177600c4 100644
--- a/docs/testing/testing-with-django.md
+++ b/docs/testing/testing-with-django.md
@@ -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
 
diff --git a/zerver/lib/test_classes.py b/zerver/lib/test_classes.py
index b7e30682de..c5776e1652 100644
--- a/zerver/lib/test_classes.py
+++ b/zerver/lib/test_classes.py
@@ -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:
diff --git a/zerver/tests/test_bots.py b/zerver/tests/test_bots.py
index dce1befd52..3b53cc017c 100644
--- a/zerver/tests/test_bots.py
+++ b/zerver/tests/test_bots.py
@@ -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")
diff --git a/zerver/tests/test_cache.py b/zerver/tests/test_cache.py
index 145d4a4598..baf1b96ba3 100644
--- a/zerver/tests/test_cache.py
+++ b/zerver/tests/test_cache.py
@@ -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):
diff --git a/zerver/tests/test_custom_profile_data.py b/zerver/tests/test_custom_profile_data.py
index 5762f9d48e..0c04d9c3b6 100644
--- a/zerver/tests/test_custom_profile_data.py
+++ b/zerver/tests/test_custom_profile_data.py
@@ -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
diff --git a/zerver/tests/test_digest.py b/zerver/tests/test_digest.py
index f0b1b96448..9d82eaed99 100644
--- a/zerver/tests/test_digest.py
+++ b/zerver/tests/test_digest.py
@@ -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)
diff --git a/zerver/tests/test_event_system.py b/zerver/tests/test_event_system.py
index ba5e834a5f..28985d0f6b 100644
--- a/zerver/tests/test_event_system.py
+++ b/zerver/tests/test_event_system.py
@@ -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):
diff --git a/zerver/tests/test_example.py b/zerver/tests/test_example.py
index 250b1d86be..ae1e36c8f9 100644
--- a/zerver/tests/test_example.py
+++ b/zerver/tests/test_example.py
@@ -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
diff --git a/zerver/tests/test_home.py b/zerver/tests/test_home.py
index 35900138bc..9275c63f12 100644
--- a/zerver/tests/test_home.py
+++ b/zerver/tests/test_home.py
@@ -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)
diff --git a/zerver/tests/test_message_dict.py b/zerver/tests/test_message_dict.py
index c034f15ef0..28995bd398 100644
--- a/zerver/tests/test_message_dict.py
+++ b/zerver/tests/test_message_dict.py
@@ -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:
diff --git a/zerver/tests/test_message_edit.py b/zerver/tests/test_message_edit.py
index cb4bd15936..3020d212cd 100644
--- a/zerver/tests/test_message_edit.py
+++ b/zerver/tests/test_message_edit.py
@@ -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")
diff --git a/zerver/tests/test_message_flags.py b/zerver/tests/test_message_flags.py
index abe24bf101..a0728116d6 100644
--- a/zerver/tests/test_message_flags.py
+++ b/zerver/tests/test_message_flags.py
@@ -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):
diff --git a/zerver/tests/test_message_send.py b/zerver/tests/test_message_send.py
index a2f464acc4..ae6c5b317a 100644
--- a/zerver/tests/test_message_send.py
+++ b/zerver/tests/test_message_send.py
@@ -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")
diff --git a/zerver/tests/test_retention.py b/zerver/tests/test_retention.py
index c72de9bd16..3532632e96 100644
--- a/zerver/tests/test_retention.py
+++ b/zerver/tests/test_retention.py
@@ -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))
diff --git a/zerver/tests/test_signup.py b/zerver/tests/test_signup.py
index 435a7940c5..233decdb00 100644
--- a/zerver/tests/test_signup.py
+++ b/zerver/tests/test_signup.py
@@ -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
diff --git a/zerver/tests/test_soft_deactivation.py b/zerver/tests/test_soft_deactivation.py
index 7a19de1cc1..12be027883 100644
--- a/zerver/tests/test_soft_deactivation.py
+++ b/zerver/tests/test_soft_deactivation.py
@@ -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()
diff --git a/zerver/tests/test_subs.py b/zerver/tests/test_subs.py
index 1319292091..c7a63a6b8e 100644
--- a/zerver/tests/test_subs.py
+++ b/zerver/tests/test_subs.py
@@ -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:
         """
diff --git a/zerver/tests/test_typing.py b/zerver/tests/test_typing.py
index 77b59dd0ab..2af6ae15b1 100644
--- a/zerver/tests/test_typing.py
+++ b/zerver/tests/test_typing.py
@@ -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"])
diff --git a/zerver/tests/test_upload.py b/zerver/tests/test_upload.py
index 35c3a07798..d88efe2cfd 100644
--- a/zerver/tests/test_upload.py
+++ b/zerver/tests/test_upload.py
@@ -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()
 
diff --git a/zerver/tests/test_users.py b/zerver/tests/test_users.py
index d866bce28e..ad31eca791 100644
--- a/zerver/tests/test_users.py
+++ b/zerver/tests/test_users.py
@@ -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)