From c73d9fb7838178bb3d01db61ef3c76c38157b0dd Mon Sep 17 00:00:00 2001
From: Ankush Menat <ankush@frappe.io>
Date: Wed, 28 Jun 2023 19:47:04 +0530
Subject: [PATCH] test: add perf test helper for counting rows read

---
 frappe/tests/test_db.py   | 21 ---------------------
 frappe/tests/test_perf.py | 20 ++++++++++++++++++++
 frappe/tests/utils.py     | 20 ++++++++++++++++++++
 3 files changed, 40 insertions(+), 21 deletions(-)

diff --git a/frappe/tests/test_db.py b/frappe/tests/test_db.py
index f026726cfa..e8afba5d96 100644
--- a/frappe/tests/test_db.py
+++ b/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())
 
diff --git a/frappe/tests/test_perf.py b/frappe/tests/test_perf.py
index 65f03c2f13..cc7d0b031d 100644
--- a/frappe/tests/test_perf.py
+++ b/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")
diff --git a/frappe/tests/utils.py b/frappe/tests/utils.py
index 6e6c72052a..762cd885b0 100644
--- a/frappe/tests/utils.py
+++ b/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):