PR Feedback 3

Author: VaggelisD

sqlmesh/core/test/__init__.py

@@ -1,18 +1,5 @@
 from __future__ import annotations
 
-import sys
-import time
-import pathlib
-import threading
-import typing as t
-import unittest
-
-
-import concurrent
-from concurrent.futures import ThreadPoolExecutor
-
-from sqlmesh.core.engine_adapter import EngineAdapter
-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 (
     ModelTestMetadata as ModelTestMetadata,
@@ -22,185 +9,6 @@
 )
 from sqlmesh.core.test.result import ModelTextTestResult as ModelTextTestResult
 from sqlmesh.core.test.runner import (
-    ModelTextTestRunner,
-    log_test_report,
+    run_model_tests as run_model_tests,
+    run_tests as run_tests,
 )
-from sqlmesh.utils import UniqueKeyDict, Verbosity
-
-if t.TYPE_CHECKING:
-    from sqlmesh.core.config.loader import C
-
-
-def run_tests(
-    model_test_metadata: list[ModelTestMetadata],
-    models: UniqueKeyDict[str, Model],
-    config: C,
-    gateway: t.Optional[str] = None,
-    dialect: str | None = None,
-    verbosity: Verbosity = Verbosity.DEFAULT,
-    preserve_fixtures: bool = False,
-    stream: t.TextIO | None = None,
-    default_catalog: str | None = None,
-    default_catalog_dialect: str = "",
-) -> ModelTextTestResult:
-    """Create a test suite of ModelTest objects and run it.
-
-    Args:
-        model_test_metadata: A list of ModelTestMetadata named tuples.
-        models: All models to use for expansion and mapping of physical locations.
-        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
-
-    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,
-    )
-
-    worker_payload = []
-
-    for metadata in model_test_metadata:
-        gateway = metadata.body.get("gateway") or default_gateway
-        test_connection = config.get_test_connection(
-            gateway, default_catalog, default_catalog_dialect
-        )
-
-        concurrent_tasks = test_connection.concurrent_tasks
-
-        from sqlmesh.core.config.connection import BaseDuckDBConnectionConfig
-
-        is_duckdb_connection = isinstance(test_connection, BaseDuckDBConnectionConfig)
-
-        engine_adapter = None
-        if is_duckdb_connection:
-            # Ensure DuckDB connections are fully isolated from each other
-            # by forcing the creation of a new adapter with SingletonConnectionPool
-            test_connection.concurrent_tasks = 1
-            engine_adapter = test_connection.create_engine_adapter(register_comments_override=False)
-            test_connection.concurrent_tasks = concurrent_tasks
-        elif gateway not in testing_adapter_by_gateway:
-            # All other engines can share connections between threads
-            testing_adapter_by_gateway[gateway] = test_connection.create_engine_adapter(
-                register_comments_override=False
-            )
-
-        engine_adapter = engine_adapter or testing_adapter_by_gateway[gateway]
-        worker_payload.append((metadata, engine_adapter))
-
-    def _run_single_test(
-        metadata: ModelTestMetadata, engine_adapter: EngineAdapter
-    ) -> ModelTextTestResult:
-        test = ModelTest.create_test(
-            body=metadata.body,
-            test_name=metadata.test_name,
-            models=models,
-            engine_adapter=engine_adapter,
-            dialect=dialect,
-            path=metadata.path,
-            default_catalog=default_catalog,
-            preserve_fixtures=preserve_fixtures,
-        )
-
-        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])
-        return result
-
-    test_results = []
-
-    workers = min(len(model_test_metadata) or 1, default_test_connection.concurrent_tasks)
-
-    start_time = time.perf_counter()
-    try:
-        with ThreadPoolExecutor(max_workers=workers) as pool:
-            futures = [
-                pool.submit(_run_single_test, metadata=metadata, engine_adapter=engine_adapter)
-                for metadata, engine_adapter in worker_payload
-            ]
-
-            for future in concurrent.futures.as_completed(futures):
-                test_results.append(future.result())
-    finally:
-        for _, engine_adapter in worker_payload:
-            if engine_adapter:
-                engine_adapter.close()
-
-    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(
-    tests: list[str],
-    models: UniqueKeyDict[str, Model],
-    config: C,
-    gateway: t.Optional[str] = None,
-    dialect: str | None = None,
-    verbosity: Verbosity = Verbosity.DEFAULT,
-    patterns: list[str] | None = None,
-    preserve_fixtures: bool = False,
-    stream: t.TextIO | None = None,
-    default_catalog: t.Optional[str] = None,
-    default_catalog_dialect: str = "",
-) -> ModelTextTestResult:
-    """Load and run tests.
-
-    Args:
-        tests: A list of tests to run, e.g. [tests/test_orders.yaml::test_single_order]
-        models: All models to use for expansion and mapping of physical locations.
-        verbosity: The verbosity level.
-        patterns: A list of patterns to match against.
-        preserve_fixtures: Preserve the fixture tables in the testing database, useful for debugging.
-    """
-    loaded_tests = []
-    for test in tests:
-        filename, test_name = test.split("::", maxsplit=1) if "::" in test else (test, "")
-        path = pathlib.Path(filename)
-
-        if test_name:
-            loaded_tests.append(load_model_test_file(path, variables=config.variables)[test_name])
-        else:
-            loaded_tests.extend(load_model_test_file(path, variables=config.variables).values())
-
-    if patterns:
-        loaded_tests = filter_tests_by_patterns(loaded_tests, patterns)
-
-    return run_tests(
-        loaded_tests,
-        models,
-        config,
-        gateway=gateway,
-        dialect=dialect,
-        verbosity=verbosity,
-        preserve_fixtures=preserve_fixtures,
-        stream=stream,
-        default_catalog=default_catalog,
-        default_catalog_dialect=default_catalog_dialect,
-    )

sqlmesh/core/test/result.py

@@ -49,7 +49,7 @@ 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)
+        self.original_err = (test, err)
         return super().addFailure(test, (exctype, value, None))  # type: ignore
 
     def addError(self, test: unittest.TestCase, err: ErrorType) -> None:
@@ -61,7 +61,7 @@ def addError(self, test: unittest.TestCase, err: ErrorType) -> None:
             test: The test case.
             err: A tuple of the form returned by sys.exc_info(), i.e., (type, value, traceback).
         """
-        self.err = (test, err)
+        self.original_err = (test, err)
         return super().addError(test, err)  # type: ignore
 
     def addSuccess(self, test: unittest.TestCase) -> None: