update tests

dagster-io · Nov 6, 2023 · a138c9f · a138c9f
1 parent fc5277d
commit a138c9f
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/python_modules/dagster/dagster/_core/definitions/time_window_partitions.py b/python_modules/dagster/dagster/_core/definitions/time_window_partitions.py
@@ -1900,6 +1900,16 @@ def __new__(
         included_time_windows: Sequence[TimeWindow] = [],
     ):
         check.opt_int_param(num_partitions, "num_partitions")
+        check.sequence_param(included_time_windows, "included_time_windows", of_type=TimeWindow)
+
+        time_windows_with_timezone = [
+            TimeWindow(
+                start=pendulum.instance(tw.start).in_tz(tz=partitions_def.timezone),
+                end=pendulum.instance(tw.end).in_tz(tz=partitions_def.timezone),
+            )
+            for tw in included_time_windows
+        ]
+
         return super(TimeWindowPartitionsSubset, cls).__new__(
             cls,
             partitions_def=partitions_def,
@@ -1908,9 +1918,7 @@ def __new__(
                 if num_partitions
                 else cls._num_partitions_from_time_windows(partitions_def, included_time_windows)
             ),
-            included_time_windows=check.sequence_param(
-                included_time_windows, "included_time_windows", of_type=TimeWindow
-            ),
+            included_time_windows=time_windows_with_timezone,
         )
 
     def get_included_time_windows(self) -> Sequence[TimeWindow]:

diff --git a/python_modules/dagster/dagster_tests/asset_defs_tests/test_partitions_subset.py b/python_modules/dagster/dagster_tests/asset_defs_tests/test_partitions_subset.py
@@ -87,7 +87,7 @@ def test_empty_subsets():
     "partitions_def",
     [
         (DailyPartitionsDefinition("2023-01-01", timezone="America/New_York")),
-        (DailyPartitionsDefinition("2023-01-01", timezone="America/New_York")),
+        (DailyPartitionsDefinition("2023-01-01")),
     ],
 )
 def test_time_window_partitions_subset_serialization_deserialization(
@@ -105,4 +105,8 @@ def test_time_window_partitions_subset_serialization_deserialization(
         time_window_partitions_def
     ).with_partition_keys(["2023-01-01"])
 
-    assert deserialize_value(serialize_value(cast(TimeWindowPartitionsSubset, subset))) == subset
+    deserialized = deserialize_value(
+        serialize_value(cast(TimeWindowPartitionsSubset, subset)), TimeWindowPartitionsSubset
+    )
+    assert deserialized == subset
+    assert deserialized.get_partition_keys() == ["2023-01-01"]