fix tests

Author: izeigerman

tests/conftest.py

@@ -470,6 +470,7 @@ def _make_function(
         dialect: t.Optional[str] = None,
         register_comments: bool = True,
         default_catalog: t.Optional[str] = None,
+        patch_get_data_objects: bool = True,
         **kwargs: t.Any,
     ) -> T:
         connection_mock = mocker.NonCallableMock()
@@ -493,6 +494,8 @@ def _make_function(
                 "sqlmesh.core.engine_adapter.mssql.MSSQLEngineAdapter.catalog_support",
                 new_callable=PropertyMock(return_value=CatalogSupport.REQUIRES_SET_CATALOG),
             )
+        if patch_get_data_objects:
+            mocker.patch.object(adapter, "_get_data_objects", return_value=[])
         return adapter
 
     return _make_function

tests/core/engine_adapter/test_athena.py

@@ -7,6 +7,7 @@
 from sqlglot import exp, parse_one
 import sqlmesh.core.dialect as d
 from sqlmesh.core.engine_adapter import AthenaEngineAdapter
+from sqlmesh.core.engine_adapter.shared import DataObject
 from sqlmesh.core.model import load_sql_based_model
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.utils.errors import SQLMeshError
@@ -288,6 +289,11 @@ def test_replace_query(adapter: AthenaEngineAdapter, mocker: MockerFixture):
         "sqlmesh.core.engine_adapter.athena.AthenaEngineAdapter._query_table_type",
         return_value="iceberg",
     )
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="", name="test", type="table")],
+    )
 
     adapter.replace_query(
         table_name="test",
@@ -304,6 +310,7 @@ def test_replace_query(adapter: AthenaEngineAdapter, mocker: MockerFixture):
     mocker.patch(
         "sqlmesh.core.engine_adapter.athena.AthenaEngineAdapter.table_exists", return_value=False
     )
+    mocker.patch.object(adapter, "_get_data_objects", return_value=[])
     adapter.cursor.execute.reset_mock()
 
     adapter.s3_warehouse_location = "s3://foo"

tests/core/engine_adapter/test_clickhouse.py

@@ -2,7 +2,7 @@
 from sqlmesh.core.engine_adapter import ClickhouseEngineAdapter
 from sqlmesh.core.model.definition import load_sql_based_model
 from sqlmesh.core.model.kind import ModelKindName
-from sqlmesh.core.engine_adapter.shared import EngineRunMode
+from sqlmesh.core.engine_adapter.shared import EngineRunMode, DataObject
 from tests.core.engine_adapter import to_sql_calls
 from sqlmesh.core.dialect import parse
 from sqlglot import exp, parse_one
@@ -573,6 +573,12 @@ def test_scd_type_2_by_time(
         make_temp_table_name(table_name, "abcd"),
     ]
 
+    mocker.patch.object(
+        adapter,
+        "get_data_objects",
+        return_value=[DataObject(schema="", name=table_name, type="table")],
+    )
+
     fetchone_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchone")
     fetchone_mock.return_value = None
 
@@ -610,7 +616,7 @@ def test_scd_type_2_by_time(
         truncate=True,
     )
 
-    assert to_sql_calls(adapter)[4] == parse_one(
+    assert to_sql_calls(adapter)[3] == parse_one(
         """
 INSERT INTO "__temp_target_abcd" ("id", "name", "price", "test_UPDATED_at", "test_valid_from", "test_valid_to")
 WITH "source" AS (
@@ -787,6 +793,12 @@ def test_scd_type_2_by_column(
         make_temp_table_name(table_name, "abcd"),
     ]
 
+    mocker.patch.object(
+        adapter,
+        "get_data_objects",
+        return_value=[DataObject(schema="", name=table_name, type="table")],
+    )
+
     fetchone_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchone")
     fetchone_mock.return_value = None
 
@@ -817,7 +829,7 @@ def test_scd_type_2_by_column(
         truncate=True,
     )
 
-    assert to_sql_calls(adapter)[4] == parse_one(
+    assert to_sql_calls(adapter)[3] == parse_one(
         """
 INSERT INTO "__temp_target_abcd" ("id", "name", "price", "test_VALID_from", "test_valid_to")
 WITH "source" AS (

tests/core/engine_adapter/test_databricks.py

@@ -8,6 +8,7 @@
 
 from sqlmesh.core import dialect as d
 from sqlmesh.core.engine_adapter import DatabricksEngineAdapter
+from sqlmesh.core.engine_adapter.shared import DataObject
 from sqlmesh.core.node import IntervalUnit
 from tests.core.engine_adapter import to_sql_calls
 
@@ -41,6 +42,11 @@ def test_replace_query_exists(mocker: MockFixture, make_mocked_engine_adapter: t
         "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog"
     )
     adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="", name="test_table", type="table")],
+    )
     adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
 
     assert to_sql_calls(adapter) == [
@@ -78,6 +84,11 @@ def test_replace_query_pandas_exists(mocker: MockFixture, make_mocked_engine_ada
         "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog"
     )
     adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="", name="test_table", type="table")],
+    )
     df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
     adapter.replace_query(
         "test_table", df, {"a": exp.DataType.build("int"), "b": exp.DataType.build("int")}

tests/core/engine_adapter/test_mssql.py

@@ -582,13 +582,16 @@ def test_merge_exists(
     ]
 
 
-def test_replace_query(make_mocked_engine_adapter: t.Callable):
+def test_replace_query(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
     adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
-    adapter.cursor.fetchone.return_value = (1,)
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="", name="test_table", type="table")],
+    )
     adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
 
     assert to_sql_calls(adapter) == [
-        """SELECT 1 FROM [INFORMATION_SCHEMA].[TABLES] WHERE [TABLE_NAME] = 'test_table';""",
         "TRUNCATE TABLE [test_table];",
         "INSERT INTO [test_table] ([a]) SELECT [a] FROM [tbl];",
     ]
@@ -605,6 +608,11 @@ def test_replace_query_pandas(
     )
 
     adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="", name="test_table", type="table")],
+    )
     adapter.cursor.fetchone.return_value = (1,)
 
     temp_table_mock = mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table")
@@ -682,7 +690,7 @@ def test_drop_schema_with_catalog(make_mocked_engine_adapter: t.Callable, mocker
 
 
 def test_get_data_objects_catalog(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
-    adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
+    adapter = make_mocked_engine_adapter(MSSQLEngineAdapter, patch_get_data_objects=False)
     original_set_current_catalog = adapter.set_current_catalog
     local_state = {}
 
@@ -912,6 +920,12 @@ def test_replace_query_strategy(adapter: MSSQLEngineAdapter, mocker: MockerFixtu
     exists_mock.return_value = True
     assert adapter.table_exists("test_table")
 
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="", name="test_table", type="table")],
+    )
+
     adapter.replace_query(
         "test_table",
         model.render_query_or_raise(),

tests/core/engine_adapter/test_redshift.py

@@ -9,6 +9,7 @@
 from sqlglot import parse_one
 
 from sqlmesh.core.engine_adapter import RedshiftEngineAdapter
+from sqlmesh.core.engine_adapter.shared import DataObject
 from sqlmesh.utils.errors import SQLMeshError
 from tests.core.engine_adapter import to_sql_calls
 
@@ -262,6 +263,11 @@ def mock_table(*args, **kwargs):
 
     mock_temp_table = mocker.MagicMock(side_effect=mock_table)
     mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table", mock_temp_table)
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="", name="test_table", type="table")],
+    )
 
     adapter.replace_query(
         table_name="test_table",

tests/core/engine_adapter/test_spark.py

@@ -10,6 +10,7 @@
 from sqlglot import parse_one
 
 from sqlmesh.core.engine_adapter import SparkEngineAdapter
+from sqlmesh.core.engine_adapter.shared import DataObject
 from sqlmesh.utils.errors import SQLMeshError
 from tests.core.engine_adapter import to_sql_calls
 import sqlmesh.core.dialect as d
@@ -102,6 +103,11 @@ def test_replace_query_table_properties_exists(
         return_value=True,
     )
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="", name="test_table", type="table")],
+    )
 
     columns_to_types = {
         "cola": exp.DataType.build("INT"),
@@ -194,6 +200,11 @@ def test_replace_query_exists(mocker: MockerFixture, make_mocked_engine_adapter:
         return_value=True,
     )
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="", name="test_table", type="table")],
+    )
     adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
 
     assert to_sql_calls(adapter) == [
@@ -239,6 +250,12 @@ def check_table_exists(table_name: exp.Table) -> bool:
         side_effect=check_table_exists,
     )
 
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="db", name="table", type="table")],
+    )
+
     adapter.replace_query(table_name, parse_one(f"SELECT col + 1 AS col FROM {table_name}"))
 
     assert to_sql_calls(adapter) == [
@@ -268,6 +285,11 @@ def test_replace_query_self_ref_exists(
 
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
     adapter.cursor.fetchone.return_value = (1,)
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="db", name="table", type="table")],
+    )
 
     table_name = "db.table"
     temp_table_id = "abcdefgh"
@@ -525,11 +547,6 @@ def test_spark_struct_complex_to_col_to_types(type_name, spark_type):
 def test_scd_type_2_by_time(
     make_mocked_engine_adapter: t.Callable, make_temp_table_name: t.Callable, mocker: MockerFixture
 ):
-    mocker.patch(
-        "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
-        return_value=False,
-    )
-
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
     adapter._default_catalog = "spark_catalog"
     adapter.spark.catalog.currentCatalog.return_value = "spark_catalog"
@@ -550,6 +567,11 @@ def check_table_exists(table_name: exp.Table) -> bool:
         "sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
         side_effect=check_table_exists,
     )
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="db", name="target", type="table")],
+    )
 
     adapter.scd_type_2_by_time(
         target_table="db.target",
@@ -981,6 +1003,11 @@ def test_replace_query_with_wap_self_reference(
     )
 
     adapter = make_mocked_engine_adapter(SparkEngineAdapter)
+    mocker.patch.object(
+        adapter,
+        "_get_data_objects",
+        return_value=[DataObject(schema="schema", name="table", type="table")],
+    )
 
     adapter.replace_query(
         "catalog.schema.table.branch_wap_12345",

tests/core/test_snapshot.py

@@ -3453,8 +3453,9 @@ def test_merge_intervals_virtual_environment_mode_dev_only_no_rebuild(make_snaps
         data_hash="different", metadata_hash="different", parent_data_hash="different"
     )
     target_snapshot.categorize_as(
-        SnapshotChangeCategory.FORWARD_ONLY
-    )  # This is a no-rebuild category
+        SnapshotChangeCategory.BREAKING,
+        forward_only=True,
+    )  # This is a no-rebuild categorization
 
     # Ensure snapshot is paused
     target_snapshot.unpaused_ts = None

tests/core/test_snapshot_evaluator.py

@@ -1262,13 +1262,11 @@ def test_migrate_missing_table(mocker: MockerFixture, make_snapshot):
 def test_migrate_view(
     mocker: MockerFixture,
     make_snapshot,
+    make_mocked_engine_adapter,
     change_category: SnapshotChangeCategory,
     forward_only: bool,
 ):
-    connection_mock = mocker.NonCallableMock()
-    cursor_mock = mocker.Mock()
-    connection_mock.cursor.return_value = cursor_mock
-    adapter = EngineAdapter(lambda: connection_mock, "")
+    adapter = make_mocked_engine_adapter(EngineAdapter)
 
     evaluator = SnapshotEvaluator(adapter)
 
@@ -1284,7 +1282,7 @@ def test_migrate_view(
 
     evaluator.migrate([snapshot], {}, deployability_index=DeployabilityIndex.none_deployable())
 
-    cursor_mock.execute.assert_has_calls(
+    adapter.cursor.execute.assert_has_calls(
         [
             call(
                 'CREATE OR REPLACE VIEW "sqlmesh__test_schema"."test_schema__test_model__1" ("c", "a") AS SELECT "c" AS "c", "a" AS "a" FROM "tbl" AS "tbl"'

tests/core/test_test.py

@@ -2153,44 +2153,6 @@ def test_test_generation_with_timestamp(tmp_path: Path) -> None:
     }
 
 
-def test_test_generation_with_decimal(tmp_path: Path, mocker: MockerFixture) -> None:
-    from decimal import Decimal
-
-    init_example_project(tmp_path, engine_type="duckdb")
-
-    config = Config(
-        default_connection=DuckDBConnectionConfig(),
-        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
-    )
-    foo_sql_file = tmp_path / "models" / "foo.sql"
-    foo_sql_file.write_text(
-        "MODEL (name sqlmesh_example.foo); SELECT dec_col FROM sqlmesh_example.bar;"
-    )
-    bar_sql_file = tmp_path / "models" / "bar.sql"
-    bar_sql_file.write_text("MODEL (name sqlmesh_example.bar); SELECT dec_col FROM external_table;")
-
-    context = Context(paths=tmp_path, config=config)
-    input_queries = {
-        '"memory"."sqlmesh_example"."bar"': "SELECT CAST(1.23 AS DECIMAL(10,2)) AS dec_col"
-    }
-
-    # DuckDB actually returns a numpy.float64, even though the value is cast into a DECIMAL,
-    # but other engines don't behave the same. E.g. BigQuery returns a proper Decimal value.
-    mocker.patch(
-        "sqlmesh.core.engine_adapter.base.EngineAdapter.fetchdf",
-        return_value=pd.DataFrame({"dec_col": [Decimal("1.23")]}),
-    )
-
-    context.create_test("sqlmesh_example.foo", input_queries=input_queries, overwrite=True)
-
-    test = load_yaml(context.path / c.TESTS / "test_foo.yaml")
-
-    assert len(test) == 1
-    assert "test_foo" in test
-    assert test["test_foo"]["inputs"] == {'"memory"."sqlmesh_example"."bar"': [{"dec_col": "1.23"}]}
-    assert test["test_foo"]["outputs"] == {"query": [{"dec_col": "1.23"}]}
-
-
 def test_test_generation_with_recursive_ctes(tmp_path: Path) -> None:
     init_example_project(tmp_path, engine_type="duckdb")
 
@@ -2247,10 +2209,6 @@ def test_test_with_gateway_specific_model(tmp_path: Path, mocker: MockerFixture)
 
     context = Context(paths=tmp_path, config=config)
     input_queries = {'"memory"."sqlmesh_example"."input_model"': "SELECT 5 AS c"}
-    mocker.patch(
-        "sqlmesh.core.engine_adapter.base.EngineAdapter.fetchdf",
-        return_value=pd.DataFrame({"c": [5]}),
-    )
 
     assert context.engine_adapter == context.engine_adapters["main"]
     with pytest.raises(