From 13304cd36dd91f9719822a972a6b43ac930a87ea Mon Sep 17 00:00:00 2001
From: Ankush Menat <ankush@frappe.io>
Date: Sat, 4 Jan 2025 10:30:30 +0530
Subject: [PATCH] ci: balance test distribution manually

Currently one test runner takes significantly longer than another. This
is entirely due to test_commands.py which needs to create new site and
do backup/restore tests etc. All of which are far far slower than other
tests.
---
 frappe/{tests => commands}/test_commands.py |  0
 frappe/commands/testing.py                  |  3 +--
 frappe/parallel_test_runner.py              | 19 +++++++++++++------
 3 files changed, 14 insertions(+), 8 deletions(-)
 rename frappe/{tests => commands}/test_commands.py (100%)

diff --git a/frappe/tests/test_commands.py b/frappe/commands/test_commands.py
similarity index 100%
rename from frappe/tests/test_commands.py
rename to frappe/commands/test_commands.py
diff --git a/frappe/commands/testing.py b/frappe/commands/testing.py
index d0e3549ed8..10dd126ec3 100644
--- a/frappe/commands/testing.py
+++ b/frappe/commands/testing.py
@@ -360,8 +360,7 @@ def run_parallel_tests(
 		║ App:            {app:<26} ║
 		║ Site:           {site:<26} ║
 		║ Build Number:   {build_number:<26} ║
-		║ Total Builds:   {total_builds:<26} ║
-		║ Tests in Build: ~{runner.total_tests:<25} ║"""
+		║ Total Builds:   {total_builds:<26} ║"""
 		if cc.with_coverage:
 			banner += """
 			║ Coverage Rep.:  {cc.outfile:<26} ║"""
diff --git a/frappe/parallel_test_runner.py b/frappe/parallel_test_runner.py
index 8c2a429a9e..c116b7d0f3 100644
--- a/frappe/parallel_test_runner.py
+++ b/frappe/parallel_test_runner.py
@@ -21,6 +21,12 @@ click_ctx = click.get_current_context(True)
 if click_ctx:
 	click_ctx.color = True
 
+TEST_WEIGHT_OVERRIDES = {
+	# XXX: command tests are significantly overweight, need a better heuristic than test count
+	# Possible better solution: stats from previous test runs.
+	"test_commands.py": 20,
+}
+
 
 class ParallelTestRunner:
 	def __init__(self, app, site, build_number=1, total_builds=1, dry_run=False):
@@ -30,7 +36,7 @@ class ParallelTestRunner:
 		self.total_builds = frappe.utils.cint(total_builds)
 		self.dry_run = dry_run
 		self.test_file_list = []
-		self.total_tests = 0
+		self.total_test_weight = 0
 		self.test_result = None
 		self.setup_test_file_list()
 
@@ -70,8 +76,7 @@ class ParallelTestRunner:
 
 	def setup_test_file_list(self):
 		self.test_file_list = self.get_test_file_list()
-		self.total_tests = sum(self.get_test_count(test) for test in self.test_file_list)
-		click.echo(f"Estimated total tests for build {self.build_number}: {self.total_tests}")
+		self.total_test_weight = sum(self.get_test_weight(test) for test in self.test_file_list)
 
 	def run_tests(self):
 		self.test_result = TestResult(stream=sys.stderr, descriptions=True, verbosity=2)
@@ -136,18 +141,20 @@ class ParallelTestRunner:
 		# Load balance based on total # of tests ~ each runner should get roughly same # of tests.
 		test_list = get_all_tests(self.app)
 
-		test_counts = [self.get_test_count(test) for test in test_list]
+		test_counts = [self.get_test_weight(test) for test in test_list]
 		test_chunks = split_by_weight(test_list, test_counts, chunk_count=self.total_builds)
 
 		return test_chunks[self.build_number - 1]
 
 	@staticmethod
-	def get_test_count(test):
+	def get_test_weight(test):
 		"""Get approximate count of tests inside a file"""
 		file_name = "/".join(test)
 
+		test_weight = TEST_WEIGHT_OVERRIDES.get(test[-1]) or 1
+
 		with open(file_name) as f:
-			test_count = f.read().count("def test_")
+			test_count = f.read().count("def test_") * test_weight
 
 		return test_count