fix tests

opensource-observer · Dec 21, 2024 · 80b67dc · 80b67dc
1 parent d55271f
commit 80b67dc
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 16 deletions.
diff --git 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
@@ -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
@@ -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