test: add perf test helper for counting rows read

frappe/tests/test_db.py

@@ -140,27 +140,6 @@ class TestDB(FrappeTestCase):
 			frappe.db.get_value("DocType", "DocField", order_by="creation desc, modified asc, name", run=0),
 		)
 
-	def test_get_value_limits(self):
-		# check both dict and list style filters
-		filters = [{"enabled": 1}, [["enabled", "=", 1]]]
-		for filter in filters:
-			self.assertEqual(1, len(frappe.db.get_values("User", filters=filter, limit=1)))
-			# count of last touched rows as per DB-API 2.0 https://peps.python.org/pep-0249/#rowcount
-			self.assertGreaterEqual(1, cint(frappe.db._cursor.rowcount))
-			self.assertEqual(2, len(frappe.db.get_values("User", filters=filter, limit=2)))
-			self.assertGreaterEqual(2, cint(frappe.db._cursor.rowcount))
-
-			# without limits length == count
-			self.assertEqual(
-				len(frappe.db.get_values("User", filters=filter)), frappe.db.count("User", filter)
-			)
-
-			frappe.db.get_value("User", filters=filter)
-			self.assertGreaterEqual(1, cint(frappe.db._cursor.rowcount))
-
-			frappe.db.exists("User", filter)
-			self.assertGreaterEqual(1, cint(frappe.db._cursor.rowcount))
-
 	def test_escape(self):
 		frappe.db.escape("香港濟生堂製藥有限公司 - IT".encode())
 

frappe/tests/test_perf.py

@@ -27,6 +27,7 @@ from frappe.model.base_document import get_controller
 from frappe.query_builder.utils import db_type_is
 from frappe.tests.test_query_builder import run_only_if
 from frappe.tests.utils import FrappeTestCase
+from frappe.utils import cint
 from frappe.website.path_resolver import PathResolver
 
 
@@ -70,6 +71,25 @@ class TestPerformance(FrappeTestCase):
 		with self.assertQueryCount(0):
 			get_controller("User")
 
+	def test_get_value_limits(self):
+		# check both dict and list style filters
+		filters = [{"enabled": 1}, [["enabled", "=", 1]]]
+		for filter in filters:
+			with self.assertRowsRead(1):
+				self.assertEqual(1, len(frappe.db.get_values("User", filters=filter, limit=1)))
+			with self.assertRowsRead(2):
+				self.assertEqual(2, len(frappe.db.get_values("User", filters=filter, limit=2)))
+
+			self.assertEqual(
+				len(frappe.db.get_values("User", filters=filter)), frappe.db.count("User", filter)
+			)
+
+			with self.assertRowsRead(1):
+				frappe.db.get_value("User", filters=filter)
+
+			with self.assertRowsRead(1):
+				frappe.db.exists("User", filter)
+
 	def test_db_value_cache(self):
 		"""Link validation if repeated should just use db.value_cache, hence no extra queries"""
 		doc = frappe.get_last_doc("User")

frappe/tests/utils.py

@@ -92,6 +92,26 @@ class FrappeTestCase(unittest.TestCase):
 		finally:
 			frappe.db.sql = orig_sql
 
+	@contextmanager
+	def assertRowsRead(self, count):
+		rows_read = 0
+
+		def _sql_with_count(*args, **kwargs):
+			nonlocal rows_read
+
+			ret = orig_sql(*args, **kwargs)
+			# count of last touched rows as per DB-API 2.0 https://peps.python.org/pep-0249/#rowcount
+			rows_read += cint(frappe.db._cursor.rowcount)
+			return ret
+
+		try:
+			orig_sql = frappe.db.sql
+			frappe.db.sql = _sql_with_count
+			yield
+			self.assertLessEqual(rows_read, count, msg="Queries read more rows than expected")
+		finally:
+			frappe.db.sql = orig_sql
+
 
 class MockedRequestTestCase(FrappeTestCase):
 	def setUp(self):