tests

python_modules/libraries/dagster-snowflake-pandas/dagster_snowflake_pandas_tests/test_snowflake_pandas_type_handler.py

@@ -36,6 +36,7 @@
 from dagster_snowflake.resources import SnowflakeConnection
 from dagster_snowflake_pandas import SnowflakePandasTypeHandler, snowflake_pandas_io_manager
 from dagster_snowflake_pandas.snowflake_pandas_type_handler import (
+    _add_missing_timezone,
     _convert_string_to_timestamp,
     _convert_timestamp_to_string,
 )
@@ -76,7 +77,9 @@ def test_handle_output():
     handler = SnowflakePandasTypeHandler()
     connection = MagicMock()
     df = DataFrame([{"col1": "a", "col2": 1}])
-    output_context = build_output_context(resource_config=resource_config)
+    output_context = build_output_context(
+        resource_config={**resource_config, "time_data_to_string": False}
+    )
 
     metadata = handler.handle_output(
         output_context,
@@ -107,7 +110,9 @@ def test_load_input():
         mock_read_sql.return_value = DataFrame([{"COL1": "a", "COL2": 1}])
 
         handler = SnowflakePandasTypeHandler()
-        input_context = build_input_context()
+        input_context = build_input_context(
+            resource_config={**resource_config, "time_data_to_string": False}
+        )
         df = handler.load_input(
             input_context,
             TableSlice(
@@ -126,8 +131,7 @@ def test_load_input():
 def test_type_conversions():
     # no timestamp data
     no_time = pandas.Series([1, 2, 3, 4, 5])
-    converted = _convert_string_to_timestamp(_convert_timestamp_to_string(no_time))
-
+    converted = _convert_string_to_timestamp(_convert_timestamp_to_string(no_time, None))
     assert (converted == no_time).all()
 
     # timestamp data
@@ -138,7 +142,7 @@ def test_type_conversions():
             pandas.Timestamp("2017-03-01T12:30:45.35"),
         ]
     )
-    time_converted = _convert_string_to_timestamp(_convert_timestamp_to_string(with_time))
+    time_converted = _convert_string_to_timestamp(_convert_timestamp_to_string(with_time, None))
 
     assert (with_time == time_converted).all()
 
@@ -148,6 +152,25 @@ def test_type_conversions():
     assert (_convert_string_to_timestamp(string_data) == string_data).all()
 
 
+def test_timezone_conversions():
+    # no timestamp data
+    no_time = pandas.Series([1, 2, 3, 4, 5])
+    converted = _add_missing_timezone(no_time, None)
+    assert (converted == no_time).all()
+
+    # timestamp data
+    with_time = pandas.Series(
+        [
+            pandas.Timestamp("2017-01-01T12:30:45.35"),
+            pandas.Timestamp("2017-02-01T12:30:45.35"),
+            pandas.Timestamp("2017-03-01T12:30:45.35"),
+        ]
+    )
+    time_converted = _add_missing_timezone(with_time, None)
+
+    assert (with_time.dt.tz_localize("UTC") == time_converted).all()
+
+
 def test_build_snowflake_pandas_io_manager():
     assert isinstance(
         build_snowflake_io_manager([SnowflakePandasTypeHandler()]), IOManagerDefinition