refactor: load balanced parallel tests without orchestrator (#18386) (#18390)

* feat: dry run in test runner for debuggin

* refactor: load balance TestRunner using # of tests

(cherry picked from commit cfee53d573)

Co-authored-by: Ankush Menat <ankush@frappe.io>

frappe/commands/utils.py

@@ -819,9 +819,16 @@ def run_tests(
 @click.option("--total-builds", help="Total number of builds", default=1)
 @click.option("--with-coverage", is_flag=True, help="Build coverage file")
 @click.option("--use-orchestrator", is_flag=True, help="Use orchestrator to run parallel tests")
+@click.option("--dry-run", is_flag=True, default=False, help="Dont actually run tests")
 @pass_context
 def run_parallel_tests(
-	context, app, build_number, total_builds, with_coverage=False, use_orchestrator=False
+	context,
+	app,
+	build_number,
+	total_builds,
+	with_coverage=False,
+	use_orchestrator=False,
+	dry_run=False,
 ):
 	with CodeCoverage(with_coverage, app):
 		site = get_site(context)
@@ -832,7 +839,13 @@ def run_parallel_tests(
 		else:
 			from frappe.parallel_test_runner import ParallelTestRunner
 
-			ParallelTestRunner(app, site=site, build_number=build_number, total_builds=total_builds)
+			ParallelTestRunner(
+				app,
+				site=site,
+				build_number=build_number,
+				total_builds=total_builds,
+				dry_run=dry_run,
+			)
 
 
 @click.command(

frappe/parallel_test_runner.py

@@ -18,11 +18,12 @@ if click_ctx:
 
 
 class ParallelTestRunner:
-	def __init__(self, app, site, build_number=1, total_builds=1):
+	def __init__(self, app, site, build_number=1, total_builds=1, dry_run=False):
 		self.app = app
 		self.site = site
 		self.build_number = frappe.utils.cint(build_number) or 1
 		self.total_builds = frappe.utils.cint(total_builds)
+		self.dry_run = dry_run
 		self.setup_test_site()
 		self.run_tests()
 
@@ -31,6 +32,9 @@ class ParallelTestRunner:
 		if not frappe.db:
 			frappe.connect()
 
+		if self.dry_run:
+			return
+
 		frappe.flags.in_test = True
 		frappe.clear_cache()
 		frappe.utils.scheduler.disable_scheduler()
@@ -64,6 +68,10 @@ class ParallelTestRunner:
 		if not file_info:
 			return
 
+		if self.dry_run:
+			print("running tests from", "/".join(file_info))
+			return
+
 		frappe.set_user("Administrator")
 		path, filename = file_info
 		module = self.get_module(path, filename)
@@ -108,12 +116,48 @@ class ParallelTestRunner:
 				sys.exit(1)
 
 	def get_test_file_list(self):
+		# Load balance based on total # of tests ~ each runner should get roughly same # of tests.
 		test_list = get_all_tests(self.app)
-		split_size = frappe.utils.ceil(len(test_list) / self.total_builds)
-		# [1,2,3,4,5,6] to [[1,2], [3,4], [4,6]] if split_size is 2
-		test_chunks = [test_list[x : x + split_size] for x in range(0, len(test_list), split_size)]
+
+		test_counts = [self.get_test_count(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):
+		"""Get approximate count of tests inside a file"""
+		file_name = "/".join(test)
+
+		with open(file_name) as f:
+			test_count = f.read().count("def test_")
+
+		return test_count
+
+
+def split_by_weight(work, weights, chunk_count):
+	"""Roughly split work by respective weight while keep ordering."""
+	expected_weight = sum(weights) // chunk_count
+
+	chunks = [[] for _ in range(chunk_count)]
+
+	chunk_no = 0
+	chunk_weight = 0
+
+	for task, weight in zip(work, weights):
+		if chunk_weight > expected_weight:
+			chunk_weight = 0
+			chunk_no += 1
+			assert chunk_no < chunk_count
+
+		chunks[chunk_no].append(task)
+		chunk_weight += weight
+
+	assert len(work) == sum(len(chunk) for chunk in chunks)
+	assert len(chunks) == chunk_count
+
+	return chunks
+
 
 class ParallelTestResult(unittest.TextTestResult):
 	def startTest(self, test):