Skip to content

Commit 1d7c7de

Browse files
izeigermanizeigerman
committed
fix tests
1 parent 224cff5 commit 1d7c7de

File tree

10 files changed

+98
-61
lines changed

10 files changed

+98
-61
lines changed

‎tests/conftest.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -470,6 +470,7 @@ def _make_function(
470470
dialect: t.Optional[str] = None,
471471
register_comments: bool = True,
472472
default_catalog: t.Optional[str] = None,
473+
patch_get_data_objects: bool = True,
473474
**kwargs: t.Any,
474475
) -> T:
475476
connection_mock = mocker.NonCallableMock()
@@ -493,6 +494,8 @@ def _make_function(
493494
"sqlmesh.core.engine_adapter.mssql.MSSQLEngineAdapter.catalog_support",
494495
new_callable=PropertyMock(return_value=CatalogSupport.REQUIRES_SET_CATALOG),
495496
)
497+
if patch_get_data_objects:
498+
mocker.patch.object(adapter, "_get_data_objects", return_value=[])
496499
return adapter
497500

498501
return _make_function

‎tests/core/engine_adapter/test_athena.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from sqlglot import exp, parse_one
88
import sqlmesh.core.dialect as d
99
from sqlmesh.core.engine_adapter import AthenaEngineAdapter
10+
from sqlmesh.core.engine_adapter.shared import DataObject
1011
from sqlmesh.core.model import load_sql_based_model
1112
from sqlmesh.core.model.definition import SqlModel
1213
from sqlmesh.utils.errors import SQLMeshError
@@ -288,6 +289,11 @@ def test_replace_query(adapter: AthenaEngineAdapter, mocker: MockerFixture):
288289
"sqlmesh.core.engine_adapter.athena.AthenaEngineAdapter._query_table_type",
289290
return_value="iceberg",
290291
)
292+
mocker.patch.object(
293+
adapter,
294+
"_get_data_objects",
295+
return_value=[DataObject(schema="", name="test", type="table")],
296+
)
291297

292298
adapter.replace_query(
293299
table_name="test",
@@ -304,6 +310,7 @@ def test_replace_query(adapter: AthenaEngineAdapter, mocker: MockerFixture):
304310
mocker.patch(
305311
"sqlmesh.core.engine_adapter.athena.AthenaEngineAdapter.table_exists", return_value=False
306312
)
313+
mocker.patch.object(adapter, "_get_data_objects", return_value=[])
307314
adapter.cursor.execute.reset_mock()
308315

309316
adapter.s3_warehouse_location = "s3://foo"

‎tests/core/engine_adapter/test_clickhouse.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from sqlmesh.core.engine_adapter import ClickhouseEngineAdapter
33
from sqlmesh.core.model.definition import load_sql_based_model
44
from sqlmesh.core.model.kind import ModelKindName
5-
from sqlmesh.core.engine_adapter.shared import EngineRunMode
5+
from sqlmesh.core.engine_adapter.shared import EngineRunMode, DataObject
66
from tests.core.engine_adapter import to_sql_calls
77
from sqlmesh.core.dialect import parse
88
from sqlglot import exp, parse_one
@@ -573,6 +573,12 @@ def test_scd_type_2_by_time(
573573
make_temp_table_name(table_name, "abcd"),
574574
]
575575

576+
mocker.patch.object(
577+
adapter,
578+
"get_data_objects",
579+
return_value=[DataObject(schema="", name=table_name, type="table")],
580+
)
581+
576582
fetchone_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchone")
577583
fetchone_mock.return_value = None
578584

@@ -610,7 +616,7 @@ def test_scd_type_2_by_time(
610616
truncate=True,
611617
)
612618

613-
assert to_sql_calls(adapter)[4] == parse_one(
619+
assert to_sql_calls(adapter)[3] == parse_one(
614620
"""
615621
INSERT INTO "__temp_target_abcd" ("id", "name", "price", "test_UPDATED_at", "test_valid_from", "test_valid_to")
616622
WITH "source" AS (
@@ -787,6 +793,12 @@ def test_scd_type_2_by_column(
787793
make_temp_table_name(table_name, "abcd"),
788794
]
789795

796+
mocker.patch.object(
797+
adapter,
798+
"get_data_objects",
799+
return_value=[DataObject(schema="", name=table_name, type="table")],
800+
)
801+
790802
fetchone_mock = mocker.patch("sqlmesh.core.engine_adapter.ClickhouseEngineAdapter.fetchone")
791803
fetchone_mock.return_value = None
792804

@@ -817,7 +829,7 @@ def test_scd_type_2_by_column(
817829
truncate=True,
818830
)
819831

820-
assert to_sql_calls(adapter)[4] == parse_one(
832+
assert to_sql_calls(adapter)[3] == parse_one(
821833
"""
822834
INSERT INTO "__temp_target_abcd" ("id", "name", "price", "test_VALID_from", "test_valid_to")
823835
WITH "source" AS (

‎tests/core/engine_adapter/test_databricks.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
from sqlmesh.core import dialect as d
1010
from sqlmesh.core.engine_adapter import DatabricksEngineAdapter
11+
from sqlmesh.core.engine_adapter.shared import DataObject
1112
from sqlmesh.core.node import IntervalUnit
1213
from tests.core.engine_adapter import to_sql_calls
1314

@@ -41,6 +42,11 @@ def test_replace_query_exists(mocker: MockFixture, make_mocked_engine_adapter: t
4142
"sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog"
4243
)
4344
adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
45+
mocker.patch.object(
46+
adapter,
47+
"_get_data_objects",
48+
return_value=[DataObject(schema="", name="test_table", type="table")],
49+
)
4450
adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
4551

4652
assert to_sql_calls(adapter) == [
@@ -78,6 +84,11 @@ def test_replace_query_pandas_exists(mocker: MockFixture, make_mocked_engine_ada
7884
"sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.set_current_catalog"
7985
)
8086
adapter = make_mocked_engine_adapter(DatabricksEngineAdapter, default_catalog="test_catalog")
87+
mocker.patch.object(
88+
adapter,
89+
"_get_data_objects",
90+
return_value=[DataObject(schema="", name="test_table", type="table")],
91+
)
8192
df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
8293
adapter.replace_query(
8394
"test_table", df, {"a": exp.DataType.build("int"), "b": exp.DataType.build("int")}

‎tests/core/engine_adapter/test_mssql.py

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -582,13 +582,16 @@ def test_merge_exists(
582582
]
583583

584584

585-
def test_replace_query(make_mocked_engine_adapter: t.Callable):
585+
def test_replace_query(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
586586
adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
587-
adapter.cursor.fetchone.return_value = (1,)
587+
mocker.patch.object(
588+
adapter,
589+
"_get_data_objects",
590+
return_value=[DataObject(schema="", name="test_table", type="table")],
591+
)
588592
adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
589593

590594
assert to_sql_calls(adapter) == [
591-
"""SELECT 1 FROM [INFORMATION_SCHEMA].[TABLES] WHERE [TABLE_NAME] = 'test_table';""",
592595
"TRUNCATE TABLE [test_table];",
593596
"INSERT INTO [test_table] ([a]) SELECT [a] FROM [tbl];",
594597
]
@@ -605,6 +608,11 @@ def test_replace_query_pandas(
605608
)
606609

607610
adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
611+
mocker.patch.object(
612+
adapter,
613+
"_get_data_objects",
614+
return_value=[DataObject(schema="", name="test_table", type="table")],
615+
)
608616
adapter.cursor.fetchone.return_value = (1,)
609617

610618
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
682690

683691

684692
def test_get_data_objects_catalog(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
685-
adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
693+
adapter = make_mocked_engine_adapter(MSSQLEngineAdapter, patch_get_data_objects=False)
686694
original_set_current_catalog = adapter.set_current_catalog
687695
local_state = {}
688696

@@ -912,6 +920,12 @@ def test_replace_query_strategy(adapter: MSSQLEngineAdapter, mocker: MockerFixtu
912920
exists_mock.return_value = True
913921
assert adapter.table_exists("test_table")
914922

923+
mocker.patch.object(
924+
adapter,
925+
"_get_data_objects",
926+
return_value=[DataObject(schema="", name="test_table", type="table")],
927+
)
928+
915929
adapter.replace_query(
916930
"test_table",
917931
model.render_query_or_raise(),

‎tests/core/engine_adapter/test_redshift.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from sqlglot import parse_one
1010

1111
from sqlmesh.core.engine_adapter import RedshiftEngineAdapter
12+
from sqlmesh.core.engine_adapter.shared import DataObject
1213
from sqlmesh.utils.errors import SQLMeshError
1314
from tests.core.engine_adapter import to_sql_calls
1415

@@ -262,6 +263,11 @@ def mock_table(*args, **kwargs):
262263

263264
mock_temp_table = mocker.MagicMock(side_effect=mock_table)
264265
mocker.patch("sqlmesh.core.engine_adapter.EngineAdapter._get_temp_table", mock_temp_table)
266+
mocker.patch.object(
267+
adapter,
268+
"_get_data_objects",
269+
return_value=[DataObject(schema="", name="test_table", type="table")],
270+
)
265271

266272
adapter.replace_query(
267273
table_name="test_table",

‎tests/core/engine_adapter/test_spark.py

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from sqlglot import parse_one
1111

1212
from sqlmesh.core.engine_adapter import SparkEngineAdapter
13+
from sqlmesh.core.engine_adapter.shared import DataObject
1314
from sqlmesh.utils.errors import SQLMeshError
1415
from tests.core.engine_adapter import to_sql_calls
1516
import sqlmesh.core.dialect as d
@@ -102,6 +103,11 @@ def test_replace_query_table_properties_exists(
102103
return_value=True,
103104
)
104105
adapter = make_mocked_engine_adapter(SparkEngineAdapter)
106+
mocker.patch.object(
107+
adapter,
108+
"_get_data_objects",
109+
return_value=[DataObject(schema="", name="test_table", type="table")],
110+
)
105111

106112
columns_to_types = {
107113
"cola": exp.DataType.build("INT"),
@@ -194,6 +200,11 @@ def test_replace_query_exists(mocker: MockerFixture, make_mocked_engine_adapter:
194200
return_value=True,
195201
)
196202
adapter = make_mocked_engine_adapter(SparkEngineAdapter)
203+
mocker.patch.object(
204+
adapter,
205+
"_get_data_objects",
206+
return_value=[DataObject(schema="", name="test_table", type="table")],
207+
)
197208
adapter.replace_query("test_table", parse_one("SELECT a FROM tbl"), {"a": "int"})
198209

199210
assert to_sql_calls(adapter) == [
@@ -239,6 +250,12 @@ def check_table_exists(table_name: exp.Table) -> bool:
239250
side_effect=check_table_exists,
240251
)
241252

253+
mocker.patch.object(
254+
adapter,
255+
"_get_data_objects",
256+
return_value=[DataObject(schema="db", name="table", type="table")],
257+
)
258+
242259
adapter.replace_query(table_name, parse_one(f"SELECT col + 1 AS col FROM {table_name}"))
243260

244261
assert to_sql_calls(adapter) == [
@@ -268,6 +285,11 @@ def test_replace_query_self_ref_exists(
268285

269286
adapter = make_mocked_engine_adapter(SparkEngineAdapter)
270287
adapter.cursor.fetchone.return_value = (1,)
288+
mocker.patch.object(
289+
adapter,
290+
"_get_data_objects",
291+
return_value=[DataObject(schema="db", name="table", type="table")],
292+
)
271293

272294
table_name = "db.table"
273295
temp_table_id = "abcdefgh"
@@ -525,11 +547,6 @@ def test_spark_struct_complex_to_col_to_types(type_name, spark_type):
525547
def test_scd_type_2_by_time(
526548
make_mocked_engine_adapter: t.Callable, make_temp_table_name: t.Callable, mocker: MockerFixture
527549
):
528-
mocker.patch(
529-
"sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
530-
return_value=False,
531-
)
532-
533550
adapter = make_mocked_engine_adapter(SparkEngineAdapter)
534551
adapter._default_catalog = "spark_catalog"
535552
adapter.spark.catalog.currentCatalog.return_value = "spark_catalog"
@@ -550,6 +567,11 @@ def check_table_exists(table_name: exp.Table) -> bool:
550567
"sqlmesh.core.engine_adapter.spark.SparkEngineAdapter.table_exists",
551568
side_effect=check_table_exists,
552569
)
570+
mocker.patch.object(
571+
adapter,
572+
"_get_data_objects",
573+
return_value=[DataObject(schema="db", name="target", type="table")],
574+
)
553575

554576
adapter.scd_type_2_by_time(
555577
target_table="db.target",
@@ -981,6 +1003,11 @@ def test_replace_query_with_wap_self_reference(
9811003
)
9821004

9831005
adapter = make_mocked_engine_adapter(SparkEngineAdapter)
1006+
mocker.patch.object(
1007+
adapter,
1008+
"_get_data_objects",
1009+
return_value=[DataObject(schema="schema", name="table", type="table")],
1010+
)
9841011

9851012
adapter.replace_query(
9861013
"catalog.schema.table.branch_wap_12345",

‎tests/core/test_snapshot.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3453,8 +3453,9 @@ def test_merge_intervals_virtual_environment_mode_dev_only_no_rebuild(make_snaps
34533453
data_hash="different", metadata_hash="different", parent_data_hash="different"
34543454
)
34553455
target_snapshot.categorize_as(
3456-
SnapshotChangeCategory.FORWARD_ONLY
3457-
) # This is a no-rebuild category
3456+
SnapshotChangeCategory.BREAKING,
3457+
forward_only=True,
3458+
) # This is a no-rebuild categorization
34583459

34593460
# Ensure snapshot is paused
34603461
target_snapshot.unpaused_ts = None

‎tests/core/test_snapshot_evaluator.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1262,13 +1262,11 @@ def test_migrate_missing_table(mocker: MockerFixture, make_snapshot):
12621262
def test_migrate_view(
12631263
mocker: MockerFixture,
12641264
make_snapshot,
1265+
make_mocked_engine_adapter,
12651266
change_category: SnapshotChangeCategory,
12661267
forward_only: bool,
12671268
):
1268-
connection_mock = mocker.NonCallableMock()
1269-
cursor_mock = mocker.Mock()
1270-
connection_mock.cursor.return_value = cursor_mock
1271-
adapter = EngineAdapter(lambda: connection_mock, "")
1269+
adapter = make_mocked_engine_adapter(EngineAdapter)
12721270

12731271
evaluator = SnapshotEvaluator(adapter)
12741272

@@ -1284,7 +1282,7 @@ def test_migrate_view(
12841282

12851283
evaluator.migrate([snapshot], {}, deployability_index=DeployabilityIndex.none_deployable())
12861284

1287-
cursor_mock.execute.assert_has_calls(
1285+
adapter.cursor.execute.assert_has_calls(
12881286
[
12891287
call(
12901288
'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

Lines changed: 0 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -2153,44 +2153,6 @@ def test_test_generation_with_timestamp(tmp_path: Path) -> None:
21532153
}
21542154

21552155

2156-
def test_test_generation_with_decimal(tmp_path: Path, mocker: MockerFixture) -> None:
2157-
from decimal import Decimal
2158-
2159-
init_example_project(tmp_path, engine_type="duckdb")
2160-
2161-
config = Config(
2162-
default_connection=DuckDBConnectionConfig(),
2163-
model_defaults=ModelDefaultsConfig(dialect="duckdb"),
2164-
)
2165-
foo_sql_file = tmp_path / "models" / "foo.sql"
2166-
foo_sql_file.write_text(
2167-
"MODEL (name sqlmesh_example.foo); SELECT dec_col FROM sqlmesh_example.bar;"
2168-
)
2169-
bar_sql_file = tmp_path / "models" / "bar.sql"
2170-
bar_sql_file.write_text("MODEL (name sqlmesh_example.bar); SELECT dec_col FROM external_table;")
2171-
2172-
context = Context(paths=tmp_path, config=config)
2173-
input_queries = {
2174-
'"memory"."sqlmesh_example"."bar"': "SELECT CAST(1.23 AS DECIMAL(10,2)) AS dec_col"
2175-
}
2176-
2177-
# DuckDB actually returns a numpy.float64, even though the value is cast into a DECIMAL,
2178-
# but other engines don't behave the same. E.g. BigQuery returns a proper Decimal value.
2179-
mocker.patch(
2180-
"sqlmesh.core.engine_adapter.base.EngineAdapter.fetchdf",
2181-
return_value=pd.DataFrame({"dec_col": [Decimal("1.23")]}),
2182-
)
2183-
2184-
context.create_test("sqlmesh_example.foo", input_queries=input_queries, overwrite=True)
2185-
2186-
test = load_yaml(context.path / c.TESTS / "test_foo.yaml")
2187-
2188-
assert len(test) == 1
2189-
assert "test_foo" in test
2190-
assert test["test_foo"]["inputs"] == {'"memory"."sqlmesh_example"."bar"': [{"dec_col": "1.23"}]}
2191-
assert test["test_foo"]["outputs"] == {"query": [{"dec_col": "1.23"}]}
2192-
2193-
21942156
def test_test_generation_with_recursive_ctes(tmp_path: Path) -> None:
21952157
init_example_project(tmp_path, engine_type="duckdb")
21962158

@@ -2247,10 +2209,6 @@ def test_test_with_gateway_specific_model(tmp_path: Path, mocker: MockerFixture)
22472209

22482210
context = Context(paths=tmp_path, config=config)
22492211
input_queries = {'"memory"."sqlmesh_example"."input_model"': "SELECT 5 AS c"}
2250-
mocker.patch(
2251-
"sqlmesh.core.engine_adapter.base.EngineAdapter.fetchdf",
2252-
return_value=pd.DataFrame({"c": [5]}),
2253-
)
22542212

22552213
assert context.engine_adapter == context.engine_adapters["main"]
22562214
with pytest.raises(

0 commit comments

Comments
 (0)