Move the patch to an autouse fixture

Author: VaggelisD

tests/__init__.py

@@ -1,26 +1,5 @@
-import typing as t
-
 from sqlmesh.core import constants as c
 from sqlmesh.core.analytics import disable_analytics
-from sqlmesh.core.config import Config, DuckDBConnectionConfig
-from sqlmesh.core.config.connection import ConnectionConfig
 
 c.MAX_FORK_WORKERS = 1
 disable_analytics()
-
-
-# Replace the implementation of Config._get_connection with a new one that always
-# returns a ConnectionConfig to simplify testing
-original_get_connection = Config.get_connection
-
-
-def _lax_get_connection(self, gateway_name: t.Optional[str] = None) -> ConnectionConfig:
-    try:
-        connection = original_get_connection(self, gateway_name)
-    except:
-        connection = DuckDBConnectionConfig()
-    return connection
-
-
-setattr(Config, "original_get_connection", original_get_connection)
-setattr(Config, "get_connection", _lax_get_connection)

tests/conftest.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+
 import datetime
 import logging
 import typing as t
 import uuid
+from contextlib import nullcontext
 from pathlib import Path
 from shutil import copytree, rmtree
 from tempfile import TemporaryDirectory
@@ -21,7 +23,8 @@
 from sqlglot.helper import ensure_list
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 
-from sqlmesh.core.config import BaseDuckDBConnectionConfig
+from sqlmesh.core.config import Config, BaseDuckDBConnectionConfig, DuckDBConnectionConfig
+from sqlmesh.core.config.connection import ConnectionConfig
 from sqlmesh.core.context import Context
 from sqlmesh.core.engine_adapter import MSSQLEngineAdapter, SparkEngineAdapter
 from sqlmesh.core.engine_adapter.base import EngineAdapter
@@ -526,3 +529,26 @@ def _make_function(table_name: str, random_id: str) -> exp.Table:
         return temp_table
 
     return _make_function
+
+
+@pytest.fixture(scope="function", autouse=True)
+def set_default_connection(request):
+    request = request.node.get_closest_marker("set_default_connection")
+    disable = request and request.kwargs.get("disable")
+
+    if disable:
+        ctx = nullcontext()
+    else:
+        original_get_connection = Config.get_connection
+
+        def _lax_get_connection(self, gateway_name: t.Optional[str] = None) -> ConnectionConfig:
+            try:
+                connection = original_get_connection(self, gateway_name)
+            except:
+                connection = DuckDBConnectionConfig()
+            return connection
+
+        ctx = mock.patch("sqlmesh.core.config.Config.get_connection", _lax_get_connection)
+
+    with ctx:
+        yield

tests/core/test_integration.py

@@ -6131,22 +6131,20 @@ def setup_senario(model_before: str, model_after: str):
     )
 
 
-@pytest.mark.isolated
+@pytest.mark.set_default_connection(disable=True)
 def test_missing_connection_config():
-    # This is the actual implementation of Config._get_connection which would be used in a real project
-    # To make testing easier, it's mocked at tests/__init__.py
-
-    with mock.patch.object(Config, "get_connection", Config.original_get_connection):
-        # Case 1: No default_connection or gateways specified should raise a ConfigError
-        with pytest.raises(ConfigError):
-            ctx = Context(config=Config())
-
-        # Case 2: No connection specified in the gateway should raise a ConfigError
-        with pytest.raises(ConfigError):
-            ctx = Context(config=Config(gateways={"default": GatewayConfig()}))
-
-        # Case 3: Specifying a default_connection or connection in the gateway should work
-        ctx = Context(config=Config(default_connection=DuckDBConnectionConfig()))
-        ctx = Context(
-            config=Config(gateways={"default": GatewayConfig(connection=DuckDBConnectionConfig())})
-        )
+    # This is testing the actual implementation of Config.get_connection
+    # To make writing tests easier, it's patched by the autouse fixture provide_sqlmesh_default_connection
+    # Case 1: No default_connection or gateways specified should raise a ConfigError
+    with pytest.raises(ConfigError):
+        ctx = Context(config=Config())
+
+    # Case 2: No connection specified in the gateway should raise a ConfigError
+    with pytest.raises(ConfigError):
+        ctx = Context(config=Config(gateways={"incorrect": GatewayConfig()}))
+
+    # Case 3: Specifying a default_connection or connection in the gateway should work
+    ctx = Context(config=Config(default_connection=DuckDBConnectionConfig()))
+    ctx = Context(
+        config=Config(gateways={"default": GatewayConfig(connection=DuckDBConnectionConfig())})
+    )