Repository: TobikoData/sqlmesh
Pull request #2553 (2024-05-02 15:31:44)

Feat!: add ability to override the engine adapter of a unit test
Addresses the 11th item in https://github.com/TobikoData/sqlmesh/issues/1637




Patch:
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index 3d44783d..1c517942 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -194,7 +194,7 @@ Usage: sqlmesh info [OPTIONS]
   Print information about a SQLMesh project.
 
   Includes counts of project models and macros and connection tests for the
-  data warehouse and test runner.
+  data warehouse.
 
 Options:
   --help  Show this message and exit.
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 8a9a3fe5..5a0195ea 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -580,7 +580,7 @@ def info(obj: Context) -> None:
     """
     Print information about a SQLMesh project.
 
-    Includes counts of project models and macros and connection tests for the data warehouse and test runner.
+    Includes counts of project models and macros and connection tests for the data warehouse.
     """
     obj.print_info()
 
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 686d67c4..5aeb0969 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -324,12 +324,9 @@ class GenericContext(BaseContext, t.Generic[C]):
         self.concurrent_tasks = concurrent_tasks or self._connection_config.concurrent_tasks
         self._engine_adapter = engine_adapter or self._connection_config.create_engine_adapter()
 
-        test_connection_config = self.config.get_test_connection(
+        self._test_connection_config = self.config.get_test_connection(
             self.gateway, self.default_catalog, default_catalog_dialect=self.engine_adapter.DIALECT
         )
-        self._test_engine_adapter = test_connection_config.create_engine_adapter(
-            register_comments_override=False
-        )
 
         self._snapshot_evaluator: t.Optional[SnapshotEvaluator] = None
 
@@ -1343,19 +1340,25 @@ class GenericContext(BaseContext, t.Generic[C]):
             for dep, query in input_queries.items()
         }
 
-        generate_test(
-            model=self.get_model(model, raise_if_missing=True),
-            input_queries=input_queries,
-            models=self._models,
-            engine_adapter=self._engine_adapter,
-            test_engine_adapter=self._test_engine_adapter,
-            project_path=self.path,
-            overwrite=overwrite,
-            variables=variables,
-            path=path,
-            name=name,
-            include_ctes=include_ctes,
-        )
+        try:
+            test_adapter = self._test_connection_config.create_engine_adapter(
+                register_comments_override=False
+            )
+            generate_test(
+                model=self.get_model(model, raise_if_missing=True),
+                input_queries=input_queries,
+                models=self._models,
+                engine_adapter=self._engine_adapter,
+                test_engine_adapter=test_adapter,
+                project_path=self.path,
+                overwrite=overwrite,
+                variables=variables,
+                path=path,
+                name=name,
+                include_ctes=include_ctes,
+            )
+        finally:
+            test_adapter.close()
 
     def test(
         self,
@@ -1372,42 +1375,44 @@ class GenericContext(BaseContext, t.Generic[C]):
         else:
             verbosity = 1
 
-        try:
-            if tests:
-                result = run_model_tests(
-                    tests=tests,
-                    models=self._models,
-                    engine_adapter=self._test_engine_adapter,
-                    dialect=self.default_dialect,
-                    verbosity=verbosity,
-                    patterns=match_patterns,
-                    preserve_fixtures=preserve_fixtures,
-                    default_catalog=self.default_catalog,
-                )
-            else:
-                test_meta = []
-
-                for path, config in self.configs.items():
-                    test_meta.extend(
-                        get_all_model_tests(
-                            path / c.TESTS,
-                            patterns=match_patterns,
-                            ignore_patterns=config.ignore_patterns,
-                        )
+        if tests:
+            result = run_model_tests(
+                tests=tests,
+                models=self._models,
+                config=self.config,
+                gateway=self.gateway,
+                dialect=self.default_dialect,
+                verbosity=verbosity,
+                patterns=match_patterns,
+                preserve_fixtures=preserve_fixtures,
+                stream=stream,
+                default_catalog=self.default_catalog,
+                default_catalog_dialect=self.engine_adapter.DIALECT,
+            )
+        else:
+            test_meta = []
+
+            for path, config in self.configs.items():
+                test_meta.extend(
+                    get_all_model_tests(
+                        path / c.TESTS,
+                        patterns=match_patterns,
+                        ignore_patterns=config.ignore_patterns,
                     )
-
-                result = run_tests(
-                    test_meta,
-                    models=self._models,
-                    engine_adapter=self._test_engine_adapter,
-                    dialect=self.default_dialect,
-                    verbosity=verbosity,
-                    preserve_fixtures=preserve_fixtures,
-                    stream=stream,
-                    default_catalog=self.default_catalog,
                 )
-        finally:
-            self._test_engine_adapter.close()
+
+            result = run_tests(
+                model_test_metadata=test_meta,
+                models=self._models,
+                config=self.config,
+                gateway=self.gateway,
+                dialect=self.default_dialect,
+                verbosity=verbosity,
+                preserve_fixtures=preserve_fixtures,
+                stream=stream,
+                default_catalog=self.default_catalog,
+                default_catalog_dialect=self.engine_adapter.DIALECT,
+            )
 
         return result
 
@@ -1558,8 +1563,6 @@ class GenericContext(BaseContext, t.Generic[C]):
         if state_connection:
             self._try_connection("state backend", state_connection.create_engine_adapter())
 
-        self._try_connection("test", self._test_engine_adapter)
-
     def close(self) -> None:
         """Releases all resources allocated by this context."""
         self.snapshot_evaluator.close()
@@ -1671,11 +1674,11 @@ class GenericContext(BaseContext, t.Generic[C]):
     def _run_plan_tests(
         self, skip_tests: bool = False
     ) -> t.Tuple[t.Optional[unittest.result.TestResult], t.Optional[str]]:
-        if self._test_engine_adapter and not skip_tests:
+        if not skip_tests:
             result, test_output = self._run_tests()
             if result.testsRun > 0:
                 self.console.log_test_results(
-                    result, test_output, self._test_engine_adapter.dialect
+                    result, test_output, self._test_connection_config._engine_adapter.DIALECT
                 )
             if not result.wasSuccessful():
                 raise PlanError(
diff --git a/sqlmesh/integrations/github/cicd/controller.py b/sqlmesh/integrations/github/cicd/controller.py
index 8c242e53..08804d95 100644
--- a/sqlmesh/integrations/github/cicd/controller.py
+++ b/sqlmesh/integrations/github/cicd/controller.py
@@ -668,7 +668,7 @@ class GithubController:
                 # Clear out console
                 self._console.consume_captured_output()
                 self._console.log_test_results(
-                    result, output, self._context._test_engine_adapter.dialect
+                    result, output, self._context._test_connection_config._engine_adapter.DIALECT
                 )
                 test_summary = self._console.consume_captured_output()
                 test_title = "Tests Passed" if result.wasSuccessful() else "Tests Failed"
diff --git a/web/server/api/endpoints/commands.py b/web/server/api/endpoints/commands.py
index 5e89d47a..2706171c 100644
--- a/web/server/api/endpoints/commands.py
+++ b/web/server/api/endpoints/commands.py
@@ -154,7 +154,7 @@ async def test(
             origin="API -> commands -> test",
         )
     context.console.log_test_results(
-        result, test_output.getvalue(), context._test_engine_adapter.dialect
+        result, test_output.getvalue(), context._test_connection_config._engine_adapter.DIALECT
     )
     return models.TestResult(
         errors=[


Test patch:
diff --git a/docs/concepts/tests.md b/docs/concepts/tests.md
index a399ad7f..7a51c834 100644
--- a/docs/concepts/tests.md
+++ b/docs/concepts/tests.md
@@ -15,6 +15,7 @@ Tests within a suite file contain the following attributes:
 * The unique name of a test
 * The name of the model targeted by this test
 * [Optional] The test's description
+* [Optional] The gateway whose `test_connection` will be used to run this test
 * Test inputs, which are defined per upstream model or external table referenced by the target model. Each test input consists of the following:
     * The name of an upstream model or external table
     * The list of rows defined as a mapping from a column name to a value associated with it
@@ -31,6 +32,7 @@ The YAML format is defined as follows:
 <unique_test_name>:
   model: <target_model_name>
   description: <description>  # Optional
+  gateway: <gateway>  # Optional
   inputs:
     <upstream_model_or_external_table_name>:
       columns:  # Optional
@@ -188,7 +190,7 @@ test_example_full_model:
         num_orders: 2
 ```
 
-### Omitting columns
+## Omitting columns
 
 Defining the complete inputs and expected outputs for wide tables, i.e. tables with many columns, can become cumbersome. Therefore, if certain columns can be safely ignored they may be omitted from any row and their value will be treated as `NULL` for that row.
 
@@ -215,7 +217,7 @@ This is useful when we can't treat the missing columns as `NULL`, but still want
 
 When `partial` is set for a _specific_ expected output, its rows need to be defined as a mapping under the `rows` key and only the columns referenced in them will be tested.
 
-### Freezing time
+## Freezing time
 
 Some models may use SQL expressions that compute datetime values at a given point in time, such as `CURRENT_TIMESTAMP`. Since these expressions are non-deterministic, it's not enough to simply specify an expected output value in order to test them.
 
@@ -364,6 +366,43 @@ Ran 2 tests in 0.024s
 OK
 ```
 
+## Using a different testing connection
+
+The testing connection can be changed for a given test. This may be useful when, e.g., the model being tested cannot be correctly transpiled to the dialect of the default testing engine.
+
+The following example demonstrates this by modifying `test_example_full_model`, so that it runs against a single-threaded local Spark process, defined as the `test_connection` of the `spark_testing` gateway in the project's `config.yaml` file:
+
+```yaml linenums="1"
+gateways:
+  local:
+    connection:
+      type: duckdb
+      database: db.db
+  spark_testing:
+    test_connection:
+      type: spark
+      config:
+        # Run Spark locally with one worker thread
+        "spark.master": "local"
+
+        # Move data under /tmp so that it is only temporarily persisted
+        "spark.sql.warehouse.dir": "/tmp/data_dir"
+        "spark.driver.extraJavaOptions": "-Dderby.system.home=/tmp/derby_dir"
+
+default_gateway: local
+
+model_defaults:
+  dialect: duckdb
+```
+
+The test would then be updated as follows:
+
+```yaml linenums="1"
+test_example_full_model:
+  gateway: spark_testing
+  # ... the other test attributes remain the same
+```
+
 ## Running tests
 
 Tests run automatically every time a new [plan](plans.md) is created, but they can also be executed on demand as described in the following sections.
diff --git a/sqlmesh/core/test/__init__.py b/sqlmesh/core/test/__init__.py
index a4787566..db5464cd 100644
--- a/sqlmesh/core/test/__init__.py
+++ b/sqlmesh/core/test/__init__.py
@@ -16,65 +16,91 @@ from sqlmesh.core.test.discovery import (
 from sqlmesh.core.test.result import ModelTextTestResult
 from sqlmesh.utils import UniqueKeyDict
 
+if t.TYPE_CHECKING:
+    from sqlmesh.core.config.loader import C
+
 
 def run_tests(
     model_test_metadata: list[ModelTestMetadata],
     models: UniqueKeyDict[str, Model],
-    engine_adapter: EngineAdapter,
+    config: C,
+    gateway: t.Optional[str] = None,
     dialect: str | None = None,
     verbosity: int = 1,
     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.
-        engine_adapter: The engine adapter to use.
         verbosity: The verbosity level.
         preserve_fixtures: Preserve the fixture tables in the testing database, useful for debugging.
     """
-    suite = unittest.TestSuite(
-        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,
+    testing_adapter_by_gateway: t.Dict[str, EngineAdapter] = {}
+    default_gateway = gateway or config.default_gateway
+
+    try:
+        tests = []
+        for metadata in model_test_metadata:
+            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
+
+            tests.append(
+                ModelTest.create_test(
+                    body=body,
+                    test_name=metadata.test_name,
+                    models=models,
+                    engine_adapter=testing_engine_adapter,
+                    dialect=dialect,
+                    path=metadata.path,
+                    default_catalog=default_catalog,
+                    preserve_fixtures=preserve_fixtures,
+                )
+            )
+
+        result = t.cast(
+            ModelTextTestResult,
+            unittest.TextTestRunner(
+                stream=stream, verbosity=verbosity, resultclass=ModelTextTestResult
+            ).run(unittest.TestSuite(tests)),
         )
-        for metadata in model_test_metadata
-    )
+    finally:
+        for testing_engine_adapter in testing_adapter_by_gateway.values():
+            testing_engine_adapter.close()
 
-    return t.cast(
-        ModelTextTestResult,
-        unittest.TextTestRunner(
-            stream=stream, verbosity=verbosity, resultclass=ModelTextTestResult
-        ).run(suite),
-    )
+    return result
 
 
 def run_model_tests(
     tests: list[str],
     models: UniqueKeyDict[str, Model],
-    engine_adapter: EngineAdapter,
+    config: C,
+    gateway: t.Optional[str] = None,
     dialect: str | None = None,
     verbosity: int = 1,
     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.
-        engine_adapter: The engine adapter to use.
         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.
@@ -95,10 +121,12 @@ def run_model_tests(
     return run_tests(
         loaded_tests,
         models,
-        engine_adapter,
+        config,
+        gateway=gateway,
         dialect=dialect,
         verbosity=verbosity,
         preserve_fixtures=preserve_fixtures,
         stream=stream,
         default_catalog=default_catalog,
+        default_catalog_dialect=default_catalog_dialect,
     )
diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
index 0e1bf116..e64ebcd7 100644
--- a/sqlmesh/core/test/definition.py
+++ b/sqlmesh/core/test/definition.py
@@ -69,16 +69,25 @@ class ModelTest(unittest.TestCase):
         self._normalized_column_name_cache: t.Dict[str, str] = {}
         self._normalized_model_name_cache: t.Dict[t.Tuple[str, bool], str] = {}
 
+        self._test_adapter_dialect = Dialect.get_or_raise(self.engine_adapter.dialect)
+
         self._validate_and_normalize_test()
 
+        if self.engine_adapter.default_catalog:
+            self._fixture_catalog: t.Optional[exp.Identifier] = exp.parse_identifier(
+                self.engine_adapter.default_catalog, dialect=self._test_adapter_dialect
+            )
+        else:
+            self._fixture_catalog = None
+
         # The test schema name is randomized to avoid concurrency issues
-        self._test_schema = f"sqlmesh_test_{random_id(short=True)}"
+        self._fixture_schema = exp.to_identifier(f"sqlmesh_test_{random_id(short=True)}")
+        self._qualified_fixture_schema = schema_(self._fixture_schema, self._fixture_catalog)
 
-        self._engine_adapter_dialect = Dialect.get_or_raise(self.engine_adapter.dialect)
-        self._transforms = self._engine_adapter_dialect.generator_class.TRANSFORMS
+        self._transforms = self._test_adapter_dialect.generator_class.TRANSFORMS
+        self._execution_time = str(self.body.get("vars", {}).get("execution_time") or "")
 
         # When execution_time is set, we mock the CURRENT_* SQL expressions so they always return it
-        self._execution_time = str(self.body.get("vars", {}).get("execution_time") or "")
         if self._execution_time:
             exec_time = exp.Literal.string(self._execution_time)
             self._transforms = {
@@ -96,6 +105,8 @@ class ModelTest(unittest.TestCase):
 
     def setUp(self) -> None:
         """Load all input tables"""
+        self.engine_adapter.create_schema(self._qualified_fixture_schema)
+
         for name, values in self.body.get("inputs", {}).items():
             all_types_are_known = False
             known_columns_to_types: t.Dict[str, exp.DataType] = {}
@@ -118,7 +129,9 @@ class ModelTest(unittest.TestCase):
                 for col, value in rows[0].items():
                     if col not in known_columns_to_types:
                         v_type = annotate_types(exp.convert(value)).type or type(value).__name__
-                        v_type = exp.maybe_parse(v_type, into=exp.DataType, dialect=self.dialect)
+                        v_type = exp.maybe_parse(
+                            v_type, into=exp.DataType, dialect=self._test_adapter_dialect
+                        )
 
                         if not type_is_known(v_type):
                             _raise_error(
@@ -131,18 +144,15 @@ class ModelTest(unittest.TestCase):
 
                         known_columns_to_types[col] = v_type
 
-            test_fixture_table = self._test_fixture_table(name)
-            self.engine_adapter.create_schema(
-                schema_(test_fixture_table.args["db"], test_fixture_table.args.get("catalog"))
-            )
-
             df = _create_df(rows, columns=known_columns_to_types)
-            self.engine_adapter.create_view(test_fixture_table, df, known_columns_to_types)
+            self.engine_adapter.create_view(
+                self._test_fixture_table(name), df, known_columns_to_types
+            )
 
     def tearDown(self) -> None:
         """Drop all fixture tables."""
         if not self.preserve_fixtures:
-            self.engine_adapter.drop_schema(self._test_schema, cascade=True)
+            self.engine_adapter.drop_schema(self._qualified_fixture_schema, cascade=True)
 
     def assert_equal(
         self,
@@ -344,7 +354,9 @@ class ModelTest(unittest.TestCase):
                     )
 
                 values["columns"] = {
-                    self._normalize_column_name(c): exp.DataType.build(t, dialect=self.dialect)
+                    self._normalize_column_name(c): exp.DataType.build(
+                        t, dialect=self._test_adapter_dialect
+                    )
                     for c, t in columns.items()
                 }
 
@@ -366,14 +378,14 @@ class ModelTest(unittest.TestCase):
     def _test_fixture_table(self, name: str) -> exp.Table:
         table = self._fixture_table_cache.get(name)
         if not table:
-            table = exp.to_table(name, dialect=self.dialect)
+            table = exp.to_table(name, dialect=self._test_adapter_dialect)
 
             # We change the table path below, so this ensures there are no name clashes
             table.this.set("this", "__".join(part.name for part in table.parts))
 
-            table.set("db", exp.to_identifier(self._test_schema))
-            if self.default_catalog:
-                table.set("catalog", exp.parse_identifier(self.default_catalog))
+            table.set("db", self._fixture_schema.copy())
+            if self._fixture_catalog:
+                table.set("catalog", self._fixture_catalog.copy())
 
             self._fixture_table_cache[name] = table
 
@@ -426,10 +438,10 @@ class SqlModelTest(ModelTest):
         # For tests we just use the model name for the table reference and we don't want to expand
         mapping = {
             name: self._test_fixture_table(name).sql()
-            for name in [
+            for name in (
                 self._normalize_model_name(name)
                 for name in self.models.keys() | self.body.get("inputs", {}).keys()
-            ]
+            )
         }
         query = self.model.render_query_or_raise(
             **self.body.get("vars", {}),
@@ -458,7 +470,7 @@ class SqlModelTest(ModelTest):
 
     def _execute(self, query: exp.Expression) -> pd.DataFrame:
         """Executes the query with the engine adapter and returns a DataFrame."""
-        with patch.dict(self._engine_adapter_dialect.generator_class.TRANSFORMS, self._transforms):
+        with patch.dict(self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms):
             return self.engine_adapter.fetchdf(query)
 
 
@@ -524,7 +536,7 @@ class PythonModelTest(ModelTest):
     def _execute_model(self) -> pd.DataFrame:
         """Executes the python model and returns a DataFrame."""
         time_ctx = freeze_time(self._execution_time) if self._execution_time else nullcontext()
-        with patch.dict(self._engine_adapter_dialect.generator_class.TRANSFORMS, self._transforms):
+        with patch.dict(self._test_adapter_dialect.generator_class.TRANSFORMS, self._transforms):
             with t.cast(AbstractContextManager, time_ctx):
                 return t.cast(
                     pd.DataFrame,
diff --git a/sqlmesh/core/test/discovery.py b/sqlmesh/core/test/discovery.py
index 60afe930..a5f2f23a 100644
--- a/sqlmesh/core/test/discovery.py
+++ b/sqlmesh/core/test/discovery.py
@@ -26,9 +26,7 @@ class ModelTestMetadata(PydanticModel):
         return self.fully_qualified_test_name.__hash__()
 
 
-def load_model_test_file(
-    path: pathlib.Path,
-) -> dict[str, ModelTestMetadata]:
+def load_model_test_file(path: pathlib.Path) -> dict[str, ModelTestMetadata]:
     """Load a single model test file.
 
     Args:
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 27d3f191..07742c12 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -12,13 +12,20 @@ from sqlglot import exp
 
 from sqlmesh.cli.example_project import init_example_project
 from sqlmesh.core import constants as c
-from sqlmesh.core.config import Config, DuckDBConnectionConfig, ModelDefaultsConfig
+from sqlmesh.core.config import (
+    Config,
+    DuckDBConnectionConfig,
+    GatewayConfig,
+    ModelDefaultsConfig,
+)
 from sqlmesh.core.context import Context
 from sqlmesh.core.dialect import parse
+from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.macros import MacroEvaluator, macro
 from sqlmesh.core.model import Model, SqlModel, load_sql_based_model, model
 from sqlmesh.core.test.definition import ModelTest, PythonModelTest, SqlModelTest
 from sqlmesh.utils.errors import ConfigError
+from sqlmesh.utils.yaml import dump as dump_yaml
 from sqlmesh.utils.yaml import load as load_yaml
 
 if t.TYPE_CHECKING:
@@ -38,7 +45,9 @@ def _create_test(
         test_name=test_name,
         model=model,
         models=context._models,
-        engine_adapter=context._test_engine_adapter,
+        engine_adapter=context._test_connection_config.create_engine_adapter(
+            register_comments_override=False
+        ),
         dialect=context.config.dialect,
         path=None,
         default_catalog=context.default_catalog,
@@ -277,9 +286,7 @@ test_foo:
 
 
 def test_row_order(sushi_context: Context, full_model_without_ctes: SqlModel) -> None:
-    model = t.cast(SqlModel, sushi_context.upsert_model(full_model_without_ctes))
-
-    # input and output rows are in different orders
+    # Input and output rows are in different orders
     body = load_yaml(
         """
 test_foo:
@@ -306,26 +313,36 @@ test_foo:
         """
     )
 
-    # model query without ORDER BY should pass unit test
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-    _check_successful_or_raise(result)
+    # Model query without ORDER BY should pass unit test
+    _check_successful_or_raise(
+        _create_test(
+            body=body,
+            test_name="test_foo",
+            model=sushi_context.upsert_model(full_model_without_ctes),
+            context=sushi_context,
+        ).run()
+    )
 
-    # model query with ORDER BY should fail unit test
     full_model_without_ctes_dict = full_model_without_ctes.dict()
     full_model_without_ctes_dict["query"] = full_model_without_ctes.query.order_by("id")  # type: ignore
     full_model_without_ctes_orderby = SqlModel(**full_model_without_ctes_dict)
 
-    model = t.cast(SqlModel, sushi_context.upsert_model(full_model_without_ctes_orderby))
-    result = _create_test(body, "test_foo", model, sushi_context).run()
-
-    expected_failure_msg = (
-        "AssertionError: Data mismatch (exp: expected, act: actual)\n\n"
-        "   id     value      ds    \n"
-        "  exp act   exp act exp act\n"
-        "0   2   1     3   2   4   3\n"
-        "1   1   2     2   3   3   4\n"
+    # Model query with ORDER BY should fail unit test
+    _check_successful_or_raise(
+        _create_test(
+            body=body,
+            test_name="test_foo",
+            model=sushi_context.upsert_model(full_model_without_ctes_orderby),
+            context=sushi_context,
+        ).run(),
+        expected_msg=(
+            "AssertionError: Data mismatch (exp: expected, act: actual)\n\n"
+            "   id     value      ds    \n"
+            "  exp act   exp act exp act\n"
+            "0   2   1     3   2   4   3\n"
+            "1   1   2     2   3   3   4\n"
+        ),
     )
-    _check_successful_or_raise(result, expected_msg=expected_failure_msg)
 
 
 def test_partial_data(sushi_context: Context) -> None:
@@ -576,8 +593,9 @@ test_child:
         d: 2020-01-01
         """
     )
+
+    mocker.patch("sqlmesh.core.test.definition.random_id", return_value="jzngz56a")
     test = _create_test(body, "test_child", child, sushi_context)
-    test._test_schema = "sqlmesh_test_jzngz56a"
 
     spy_execute = mocker.spy(test.engine_adapter, "_execute")
     _check_successful_or_raise(test.run())
@@ -919,6 +937,7 @@ test_foo:
 
 
 def test_freeze_time(mocker: MockerFixture) -> None:
+    mocker.patch("sqlmesh.core.test.definition.random_id", return_value="jzngz56a")
     test = _create_test(
         body=load_yaml(
             """
@@ -939,20 +958,20 @@ test_foo:
         ),
         context=Context(config=Config(model_defaults=ModelDefaultsConfig(dialect="duckdb"))),
     )
-    test._test_schema = f"sqlmesh_test_jzngz56a"
 
     spy_execute = mocker.spy(test.engine_adapter, "_execute")
     _check_successful_or_raise(test.run())
 
     spy_execute.assert_has_calls(
         [
+            call('CREATE SCHEMA IF NOT EXISTS "memory"."sqlmesh_test_jzngz56a"'),
             call(
                 "SELECT "
                 """CAST('2023-01-01 12:05:03+00:00' AS DATE) AS "cur_date", """
                 """CAST('2023-01-01 12:05:03+00:00' AS TIME) AS "cur_time", """
                 '''CAST('2023-01-01 12:05:03+00:00' AS TIMESTAMP) AS "cur_timestamp"'''
             ),
-            call('DROP SCHEMA IF EXISTS "sqlmesh_test_jzngz56a" CASCADE'),
+            call('DROP SCHEMA IF EXISTS "memory"."sqlmesh_test_jzngz56a" CASCADE'),
         ]
     )
 
@@ -994,7 +1013,8 @@ def test_successes(sushi_context: Context) -> None:
     assert "test_customer_revenue_by_day" in successful_tests
 
 
-def test_create_external_model_fixture(sushi_context: Context) -> None:
+def test_create_external_model_fixture(sushi_context: Context, mocker: MockerFixture) -> None:
+    mocker.patch("sqlmesh.core.test.definition.random_id", return_value="jzngz56a")
     test = _create_test(
         body=load_yaml(
             """
@@ -1012,7 +1032,6 @@ test_foo:
         model=_create_model("SELECT x FROM c.db.external"),
         context=sushi_context,
     )
-    test._test_schema = f"sqlmesh_test_jzngz56a"
     _check_successful_or_raise(test.run())
 
     assert len(test._fixture_table_cache) == len(sushi_context.models) + 1
@@ -1044,6 +1063,56 @@ test_foo:
     )
 
 
+def test_gateway(copy_to_temp_path: t.Callable, mocker: MockerFixture) -> None:
+    path = Path(copy_to_temp_path("examples/sushi")[0])
+    db_db_path = str(path / "db.db")
+    test_db_path = str(path / "test.db")
+
+    config = Config(
+        gateways={
+            "main": GatewayConfig(connection=DuckDBConnectionConfig(database=db_db_path)),
+            "test": GatewayConfig(test_connection=DuckDBConnectionConfig(database=test_db_path)),
+        },
+        default_gateway="main",
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+    )
+
+    context = Context(paths=path, config=config)
+
+    test_path = path / c.TESTS / "test_customer_revenue_by_day.yaml"
+    test_dict = load_yaml(test_path)
+    test_dict["test_customer_revenue_by_day"]["gateway"] = "test"
+
+    with open(test_path, "w", encoding="utf-8") as file:
+        dump_yaml(test_dict, file)
+
+    spy_execute = mocker.spy(EngineAdapter, "_execute")
+    mocker.patch("sqlmesh.core.test.definition.random_id", return_value="jzngz56a")
+
+    result = context.test(tests=[f"{test_path}::test_customer_revenue_by_day"])
+    _check_successful_or_raise(result)
+
+    expected_view_sql = (
+        'CREATE OR REPLACE VIEW "test"."sqlmesh_test_jzngz56a"."db__sushi__orders" '
+        '("id", "customer_id", "waiter_id", "start_ts", "end_ts", "event_date") AS '
+        "SELECT "
+        'CAST("id" AS INT) AS "id", '
+        'CAST("customer_id" AS INT) AS "customer_id", '
+        'CAST("waiter_id" AS INT) AS "waiter_id", '
+        'CAST("start_ts" AS INT) AS "start_ts", '
+        'CAST("end_ts" AS INT) AS "end_ts", '
+        'CAST("event_date" AS DATE) AS "event_date" '
+        "FROM (VALUES "
+        "(1, 1, 1, 1641002340, 1641004140, CAST('2022-01-01' AS DATE)), "
+        "(2, 1, 2, 1641007740, 1641009540, CAST('2022-01-01' AS DATE))) "
+        'AS "t"("id", "customer_id", "waiter_id", "start_ts", "end_ts", "event_date")'
+    )
+    test_adapter = t.cast(ModelTest, result.successes[0]).engine_adapter
+    assert call(test_adapter, expected_view_sql) in spy_execute.mock_calls
+
+    _check_successful_or_raise(context.test())
+
+
 def test_test_generation(tmp_path: Path) -> None:
     init_example_project(tmp_path, dialect="duckdb")
 
@@ -1099,8 +1168,7 @@ def test_test_generation(tmp_path: Path) -> None:
     assert "ctes" in test["test_full_model"]["outputs"]
     assert "cte" in test["test_full_model"]["outputs"]["ctes"]
 
-    result = context.test()
-    _check_successful_or_raise(result)
+    _check_successful_or_raise(context.test())
 
     context.create_test(
         "sqlmesh_example.full_model",
diff --git a/tests/integrations/jupyter/test_magics.py b/tests/integrations/jupyter/test_magics.py
index 5ab862c8..662729bf 100644
--- a/tests/integrations/jupyter/test_magics.py
+++ b/tests/integrations/jupyter/test_magics.py
@@ -510,12 +510,11 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_
 
     assert not output.stdout
     assert not output.stderr
-    assert len(output.outputs) == 4
+    assert len(output.outputs) == 3
     assert convert_all_html_output_to_text(output) == [
         "Models: 16",
         "Macros: 5",
         "Data warehouse connection succeeded",
-        "Test connection succeeded",
     ]
     assert get_all_html_output(output) == [
         str(
@@ -561,20 +560,6 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_
                 ),
             )
         ),
-        str(
-            h(
-                "pre",
-                {"style": RICH_PRE_STYLE},
-                "Test connection "
-                + str(
-                    h(
-                        "span",
-                        {"style": SUCCESS_STYLE},
-                        "succeeded",
-                    )
-                ),
-            )
-        ),
     ]