Feat: Make model tests concurrent

Author: VaggelisD

sqlmesh/core/test/__init__.py

@@ -1,10 +1,16 @@
 from __future__ import annotations
 
+import sys
+import time
 import pathlib
+import threading
 import typing as t
 import unittest
 
-from sqlmesh.core.engine_adapter import EngineAdapter
+
+import concurrent
+from concurrent.futures import ThreadPoolExecutor
+
 from sqlmesh.core.model import Model
 from sqlmesh.core.test.definition import ModelTest as ModelTest, generate_test as generate_test
 from sqlmesh.core.test.discovery import (
@@ -20,6 +26,66 @@
     from sqlmesh.core.config.loader import C
 
 
+class ModelTextTestRunner(unittest.TextTestRunner):
+    def __init__(
+        self,
+        **kwargs: t.Any,
+    ) -> None:
+        # StringIO is used to capture the output of the tests since we'll
+        # run them in parallel and we don't want to mix the output streams
+        from io import StringIO
+
+        super().__init__(
+            stream=StringIO(),
+            resultclass=ModelTextTestResult,
+            **kwargs,
+        )
+
+
+def log_test_report(results: ModelTextTestResult, test_duration: float) -> None:
+    # Aggregate parallel test run results
+    tests_run = results.testsRun
+    errors = results.errors
+    failures = results.failures
+    skipped = results.skipped
+
+    is_success = not (errors or failures)
+
+    # Compute test info
+    infos = []
+    if failures:
+        infos.append(f"failures={len(failures)}")
+    if errors:
+        infos.append(f"errors={len(errors)}")
+    if skipped:
+        infos.append(f"skipped={skipped}")
+
+    # Report test errors
+    stream = results.stream
+
+    stream.write("\n")
+
+    if errors or failures:
+        stream.writeln(unittest.TextTestResult.separator1)
+        for failure in failures:
+            stream.writeln(f"FAIL: {failure[0]}")
+
+        stream.writeln(unittest.TextTestResult.separator2)
+        for error in errors:
+            stream.writeln(error[1])
+        for failure in failures:
+            stream.writeln(failure[1])
+
+    # Test report
+    stream.writeln(unittest.TextTestResult.separator2)
+    stream.writeln(
+        f'Ran {tests_run} {"tests" if tests_run > 1 else "test"} in {test_duration:.3f}s \n'
+    )
+    stream.write(
+        f'{"OK" if is_success else "FAILED"}{" (" + ", ".join(infos) + ")" if infos else ""}'
+    )
+
+
 def run_tests(
     model_test_metadata: list[ModelTestMetadata],
     models: UniqueKeyDict[str, Model],
@@ -40,22 +106,35 @@ def run_tests(
         verbosity: The verbosity level.
         preserve_fixtures: Preserve the fixture tables in the testing database, useful for debugging.
     """
-    testing_adapter_by_gateway: t.Dict[str, EngineAdapter] = {}
     default_gateway = gateway or config.default_gateway_name
 
-    try:
-        tests = []
-        for metadata in model_test_metadata:
+    default_test_connection = config.get_test_connection(
+        gateway_name=default_gateway,
+        default_catalog=default_catalog,
+        default_catalog_dialect=default_catalog_dialect,
+    )
+
+    lock = threading.Lock()
+
+    combined_results = ModelTextTestResult(
+        stream=unittest.runner._WritelnDecorator(stream or sys.stderr),  # type: ignore
+        verbosity=2 if verbosity >= Verbosity.VERBOSE else 1,
+        descriptions=None,
+    )
+
+    def _run_single_test(metadata: ModelTestMetadata) -> ModelTextTestResult:
+        testing_engine_adapter = None
+
+        try:
             body = metadata.body
             gateway = body.get("gateway") or default_gateway
-            testing_engine_adapter = testing_adapter_by_gateway.get(gateway)
-            if not testing_engine_adapter:
-                testing_engine_adapter = config.get_test_connection(
-                    gateway,
-                    default_catalog,
-                    default_catalog_dialect,
-                ).create_engine_adapter(register_comments_override=False)
-                testing_adapter_by_gateway[gateway] = testing_engine_adapter
+
+            # Create new connection for each test to avoid concurrency issues
+            testing_engine_adapter = config.get_test_connection(
+                gateway,
+                default_catalog,
+                default_catalog_dialect,
+            ).create_engine_adapter(register_comments_override=False)
 
             test = ModelTest.create_test(
                 body=body,
@@ -67,22 +146,49 @@ def run_tests(
                 default_catalog=default_catalog,
                 preserve_fixtures=preserve_fixtures,
             )
-            if test:
-                tests.append(test)
-
-        result = t.cast(
-            ModelTextTestResult,
-            unittest.TextTestRunner(
-                stream=stream,
-                verbosity=2 if verbosity >= Verbosity.VERBOSE else 1,
-                resultclass=ModelTextTestResult,
-            ).run(unittest.TestSuite(tests)),
-        )
-    finally:
-        for testing_engine_adapter in testing_adapter_by_gateway.values():
-            testing_engine_adapter.close()
 
-    return result
+            result = t.cast(
+                ModelTextTestResult,
+                ModelTextTestRunner().run(t.cast(unittest.TestCase, test)),
+            )
+
+            with lock:
+                if result.successes:
+                    combined_results.addSuccess(result.successes[0])
+                elif result.errors:
+                    combined_results.addError(result.err[0], result.err[1])
+                elif result.failures:
+                    combined_results.addFailure(result.err[0], result.err[1])
+                elif result.skipped:
+                    skipped_args = result.skipped[0]
+                    combined_results.addSkip(skipped_args[0], skipped_args[1])
+
+        finally:
+            if testing_engine_adapter:
+                testing_engine_adapter.close()
+
+        return result
+
+    test_results = []
+
+    workers = min(len(model_test_metadata) or 1, default_test_connection.concurrent_tasks)
+
+    start_time = time.perf_counter()
+    with ThreadPoolExecutor(max_workers=workers) as pool:
+        futures = [
+            pool.submit(_run_single_test, metadata=metadata) for metadata in model_test_metadata
+        ]
+
+        for future in concurrent.futures.as_completed(futures):
+            test_results.append(future.result())
+
+    end_time = time.perf_counter()
+
+    combined_results.testsRun = len(test_results)
+
+    log_test_report(combined_results, test_duration=end_time - start_time)
+
+    return combined_results
 
 
 def run_model_tests(

sqlmesh/core/test/result.py

@@ -49,8 +49,21 @@ def addFailure(self, test: unittest.TestCase, err: ErrorType) -> None:
             err: A tuple of the form returned by sys.exc_info(), i.e., (type, value, traceback).
         """
         exctype, value, tb = err
+        self.err = (test, err)
         return super().addFailure(test, (exctype, value, None))  # type: ignore
 
+    def addError(self, test: unittest.TestCase, err: ErrorType) -> None:
+        """Called when the test case test signals an error.
+
+        The traceback is suppressed because it is redundant and not useful.
+
+        Args:
+            test: The test case.
+            err: A tuple of the form returned by sys.exc_info(), i.e., (type, value, traceback).
+        """
+        self.err = (test, err)
+        return super().addError(test, err)  # type: ignore
+
     def addSuccess(self, test: unittest.TestCase) -> None:
         """Called when the test case test succeeds.