diff --git a/warehouse/metrics_tools/utils/test_tables.py b/warehouse/metrics_tools/utils/test_tables.py index 4db0b6a81..7415465cb 100644 --- a/warehouse/metrics_tools/utils/test_tables.py +++ b/warehouse/metrics_tools/utils/test_tables.py @@ -7,7 +7,7 @@ def test_create_dependent_tables_map(): mock = MagicMock(name="context") - mock.table.return_value = "test_table" + mock.resolve_table.return_value = "test_table" actual_tables_map = create_dependent_tables_map(mock, "select * from foo") expected_tables_map = { @@ -125,7 +125,7 @@ def test_create_dependent_tables_map_parameterized( input: str, expected: t.Dict[str, str] ): mock = MagicMock(name="context") - mock.table.return_value = "test_table" + mock.resolve_table.return_value = "test_table" actual_tables_map = create_dependent_tables_map(mock, input) assert actual_tables_map == expected diff --git a/warehouse/oso_dagster/cbt/transforms/test_transforms.py b/warehouse/oso_dagster/cbt/transforms/test_transforms.py index 5aa4dac7d..ff58af0bc 100644 --- a/warehouse/oso_dagster/cbt/transforms/test_transforms.py +++ b/warehouse/oso_dagster/cbt/transforms/test_transforms.py @@ -39,18 +39,18 @@ def db(): { "time_and_id": [ ["time", "id"], - [arrow.get("2024-01-01").datetime, 1], - [arrow.get("2024-02-01").datetime, 2], - [arrow.get("2024-03-01").datetime, 3], - [arrow.get("2024-04-01").datetime, 4], + [arrow.get("2024-01-01", tzinfo="UTC").datetime, 1], + [arrow.get("2024-02-01", tzinfo="UTC").datetime, 2], + [arrow.get("2024-03-01", tzinfo="UTC").datetime, 3], + [arrow.get("2024-04-01", tzinfo="UTC").datetime, 4], ], "time_id_and_name": [ ["time", "id", "name"], - [arrow.get("2024-01-01").datetime, 1, "alpha"], - [arrow.get("2024-02-01").datetime, 2, "bravo"], - [arrow.get("2024-03-01").datetime, 3, "charlie"], - [arrow.get("2024-04-01").datetime, 4, "delta"], - [arrow.get("2024-02-01").datetime, 5, "foxtrot"], + [arrow.get("2024-01-01", tzinfo="UTC").datetime, 1, "alpha"], + [arrow.get("2024-02-01", tzinfo="UTC").datetime, 2, "bravo"], + [arrow.get("2024-03-01", tzinfo="UTC").datetime, 3, "charlie"], + [arrow.get("2024-04-01", tzinfo="UTC").datetime, 4, "delta"], + [arrow.get("2024-02-01", tzinfo="UTC").datetime, 5, "foxtrot"], ], } ) @@ -63,12 +63,15 @@ def db(): [ ( SELECT_NO_CTE, - dict(start=arrow.get("2024-02-01"), end=arrow.get("2024-04-01")), + dict( + start=arrow.get("2024-02-01", tzinfo="UTC"), + end=arrow.get("2024-04-01", tzinfo="UTC"), + ), 2, ), ( SELECT_WITH_CTE, - dict(start=arrow.get("2024-02-01")), + dict(start=arrow.get("2024-02-01", tzinfo="UTC")), 3, ), ], @@ -93,7 +96,10 @@ def test_time_constrain_succeed( ( SELECT_WITH_JOIN, "time_id_and_name", - dict(start=arrow.get("2024-02-01"), end=arrow.get("2024-04-01")), + dict( + start=arrow.get("2024-02-01", tzinfo="UTC"), + end=arrow.get("2024-04-01", tzinfo="UTC"), + ), 2, ), ], diff --git a/warehouse/oso_dagster/utils/testing/duckdb.py b/warehouse/oso_dagster/utils/testing/duckdb.py index 06d8e1cf5..353997c50 100644 --- a/warehouse/oso_dagster/utils/testing/duckdb.py +++ b/warehouse/oso_dagster/utils/testing/duckdb.py @@ -1,7 +1,7 @@ -from typing import Optional, Dict, List +from typing import Dict, List, Optional -import pandas import duckdb +import pandas class DuckDbFixture: @@ -12,6 +12,7 @@ def setup(cls, tables: Optional[Dict[str, List[List]]] = None): tables = tables or dict() db = duckdb.connect() fixture = cls(db) + fixture.set_timezone() fixture.add_tables(tables) fixture.set_memory_limit("2MB") return fixture @@ -36,6 +37,9 @@ def add_tables(self, tables: Dict[str, List]): def set_memory_limit(self, limit: str): self._db.sql(f"SET memory_limit = '{limit}';") + def set_timezone(self): + self._db.sql("SET TimeZone = 'UTC';") + @property def db(self): return self._db