diff --git a/tools/testing/kunit/kunit.py b/tools/testing/kunit/kunit.py
index 9bc0324..3c5ef46 100755
--- a/tools/testing/kunit/kunit.py
+++ b/tools/testing/kunit/kunit.py
@@ -22,34 +22,42 @@
 		    metavar='timeout')
 
 cli_args = parser.parse_args()
-linux = kunit_kernel.LinuxSourceTree()
 
-config_start = time.time()
-success = linux.build_reconfig()
-config_end = time.time()
-if not success:
-	quit()
+def main(linux):
 
-kunit_parser.print_with_timestamp('Building KUnit Kernel ...')
+	config_start = time.time()
+	success = linux.build_reconfig()
+	config_end = time.time()
+	if not success:
+		return
 
-build_start = time.time()
-success = linux.build_um_kernel()
-build_end = time.time()
-if not success:
-	quit()
+	print(kunit_parser.timestamp('Building KUnit Kernel ...'))
 
-kunit_parser.print_with_timestamp('Starting KUnit Kernel ...')
-test_start = time.time()
+	build_start = time.time()
+	success = linux.build_um_kernel()
+	build_end = time.time()
+	if not success:
+		return
 
-if cli_args.raw_output:
-	kunit_parser.raw_output(linux.run_kernel(timeout=cli_args.timeout))
-else:
-	kunit_parser.parse_run_tests(linux.run_kernel(timeout=cli_args.timeout))
+	print(kunit_parser.timestamp('Starting KUnit Kernel ...'))
+	test_start = time.time()
 
-test_end = time.time()
+	if cli_args.raw_output:
+		kunit_parser.raw_output(
+			linux.run_kernel(timeout=cli_args.timeout))
+	else:
+		for line in kunit_parser.parse_run_tests(
+			kunit_parser.isolate_kunit_output(
+				linux.run_kernel(timeout=cli_args.timeout))):
+			print(line)
 
-kunit_parser.print_with_timestamp((
-	"Elapsed time: %.3fs total, %.3fs configuring, %.3fs " +
-	"building, %.3fs running.\n") % (test_end - config_start,
-	config_end - config_start, build_end - build_start,
-	test_end - test_start))
+	test_end = time.time()
+
+	print(kunit_parser.timestamp((
+		'Elapsed time: %.3fs total, %.3fs configuring, %.3fs ' +
+		'building, %.3fs running.\n') % (test_end - config_start,
+		config_end - config_start, build_end - build_start,
+		test_end - test_start)))
+
+if __name__ == '__main__':
+	main(kunit_kernel.LinuxSourceTree())
\ No newline at end of file
diff --git a/tools/testing/kunit/kunit_parser.py b/tools/testing/kunit/kunit_parser.py
index 0ec2299..67ffb7b 100644
--- a/tools/testing/kunit/kunit_parser.py
+++ b/tools/testing/kunit/kunit_parser.py
@@ -37,12 +37,12 @@
 def green(text):
 	return '\033[1;32m' + text + RESET
 
-def print_with_timestamp(message):
-	print('[%s] %s' % (datetime.now().strftime('%H:%M:%S'), message))
+def timestamp(message):
+	return '[%s] %s' % (datetime.now().strftime('%H:%M:%S'), message)
 
-def print_log(log):
+def timestamp_log(log):
 	for m in log:
-		print_with_timestamp(m)
+		yield timestamp(m)
 
 def parse_run_tests(kernel_output):
 	test_case_output = re.compile('^kunit .*?: (.*)$')
@@ -68,18 +68,18 @@
 		log.clear()
 		total_tests.add(get_test_name(match))
 
-	print_with_timestamp(DIVIDER)
+	yield timestamp(DIVIDER)
 	try:
-		for line in isolate_kunit_output(kernel_output):
+		for line in kernel_output:
 			# Ignore module output:
 			if (test_module_success.match(line) or
 			test_module_fail.match(line)):
-				print_with_timestamp(DIVIDER)
+				yield timestamp(DIVIDER)
 				continue
 
 			match = re.match(test_case_success, line)
 			if match:
-				print_with_timestamp(green("[PASSED] ") +
+				yield timestamp(green("[PASSED] ") +
 						get_test_name(match))
 				end_one_test(match, current_case_log)
 				continue
@@ -89,20 +89,20 @@
 			# want to show and count it once.
 			if match and get_test_name(match) not in crashed_tests:
 				failed_tests.add(get_test_name(match))
-				print_with_timestamp(red("[FAILED] " +
+				yield timestamp(red("[FAILED] " +
 							get_test_name(match)))
-				print_log(map(yellow, current_case_log))
-				print_with_timestamp("")
+				yield from timestamp_log(map(yellow, current_case_log))
+				yield timestamp("")
 				end_one_test(match, current_case_log)
 				continue
 
 			match = re.match(test_case_crash, line)
 			if match:
 				crashed_tests.add(get_test_name(match))
-				print_with_timestamp(yellow("[CRASH] " +
+				yield timestamp(yellow("[CRASH] " +
 							get_test_name(match)))
-				print_log(current_case_log)
-				print_with_timestamp("")
+				yield from timestamp_log(current_case_log)
+				yield timestamp("")
 				end_one_test(match, current_case_log)
 				continue
 
@@ -114,21 +114,21 @@
 				current_case_log.append(line)
 	except KernelCrashException:
 		did_kernel_crash = True
-		print_with_timestamp(
+		yield timestamp(
 			red("The KUnit kernel crashed unexpectedly and was " +
 			    "unable to finish running tests!"))
-		print_with_timestamp(red("These are the logs from the most " +
+		yield timestamp(red("These are the logs from the most " +
 					 "recently running test:"))
-		print_with_timestamp(DIVIDER)
-		print_log(current_case_log)
-		print_with_timestamp(DIVIDER)
+		yield timestamp(DIVIDER)
+		yield from timestamp_log(current_case_log)
+		yield timestamp(DIVIDER)
 
 	fmt = green if (len(failed_tests) + len(crashed_tests) == 0
 			and not did_kernel_crash) else red
 	message = ("Before the crash:" if did_kernel_crash else
 		   "Testing complete.")
 
-	print_with_timestamp(
+	yield timestamp(
 		fmt(message + " %d tests run. %d failed. %d crashed." %
 		    (len(total_tests), len(failed_tests), len(crashed_tests))))