From c6f515a49818734617c745696a370f2dd6e45ce4 Mon Sep 17 00:00:00 2001 From: JamieDeMaria Date: Wed, 22 Mar 2023 16:06:37 -0400 Subject: [PATCH] tests --- .../test_snowflake_pandas_type_handler.py | 33 ++++++++++++++++--- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/python_modules/libraries/dagster-snowflake-pandas/dagster_snowflake_pandas_tests/test_snowflake_pandas_type_handler.py b/python_modules/libraries/dagster-snowflake-pandas/dagster_snowflake_pandas_tests/test_snowflake_pandas_type_handler.py index 5137981eb66ac..b0657f230b85a 100644 --- a/python_modules/libraries/dagster-snowflake-pandas/dagster_snowflake_pandas_tests/test_snowflake_pandas_type_handler.py +++ b/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