added integration tests for timezone aware pandas dataframe and series

Author: dimosped

arctic/store/_pandas_ndarray_store.py

@@ -1,6 +1,7 @@
 import ast
 import logging
 
+from arctic._util import NP_OBJECT_DTYPE
 from bson.binary import Binary
 from pandas import DataFrame, Series, Panel
 import numpy as np

tests/integration/store/test_version_store.py

@@ -3,7 +3,7 @@
 import struct
 from datetime import datetime as dt, timedelta as dtd
 import pandas as pd
-from arctic import VERSION_STORE
+from arctic import VERSION_STORE, PandasDataFrameStore, PandasSeriesStore
 from pandas.util.testing import assert_frame_equal, assert_series_equal
 from pymongo.errors import OperationFailure
 from pymongo.server_type import SERVER_TYPE
@@ -1587,3 +1587,23 @@ def test_write_series_with_some_objects(library, input_series):
     library.write(symbol='symX', data=input_series)
     read_data = library.read(symbol='symX').data
     assert_series_equal(input_series, read_data)
+
+
+def test_can_write_tz_aware_data_df(library):
+    mydf = _mixed_test_data()['index_tz_aware'][0]
+    library.write(symbol='symTz', data=mydf)
+    read_data = library.read(symbol='symTz').data
+    # Arctic converts by default the data to UTC, convert back
+    read_data.colB = read_data.colB.dt.tz_localize('UTC').dt.tz_convert(read_data.index.tzinfo)
+    assert library._versions.find_one({'symbol': 'symTz'})['type'] == PandasDataFrameStore.TYPE
+    assert_frame_equal(mydf, read_data)
+
+
+def test_can_write_tz_aware_data_series(library):
+    myseries = _mixed_test_data()['index_tz_aware'][0]['colB']
+    library.write(symbol='symTzSer', data=myseries)
+    read_data = library.read(symbol='symTzSer').data
+    # Arctic converts by default the data to UTC, convert back
+    read_data = read_data.dt.tz_localize('UTC').dt.tz_convert(read_data.index.tzinfo)
+    assert library._versions.find_one({'symbol': 'symTzSer'})['type'] == PandasSeriesStore.TYPE
+    assert_series_equal(myseries, read_data)

tests/unit/serialization/serialization_test_data.py

@@ -105,6 +105,12 @@ def _new_np_nd_array(val):
         multi_column_with_some_objects = multi_column_no_multiindex.copy()
         multi_column_with_some_objects.iloc[1:, 1:2] = 'Convert this columnt dtype to object'
 
+        # Index with timezone-aware datetime
+        index_tz_aware = pd.DataFrame(data={'colA': range(10),
+                                            'colB': pd.date_range('20130101', periods=10, tz='US/Eastern')},
+                                      index=pd.date_range('20130101', periods=10, tz='US/Eastern'))
+        index_tz_aware.index.name = 'index'
+
         _TEST_DATA = {
             'onerow': (onerow_ts, df_serializer.serialize(onerow_ts),
                        df_serializer.can_convert_to_records_without_objects(small_ts, 'symA')),
@@ -156,7 +162,9 @@ def _new_np_nd_array(val):
                 df_serializer.can_convert_to_records_without_objects(multi_column_with_some_objects, 'symA')),
             'n_dimensional_df': (n_dimensional_df, Exception, None),
             'mixed_dtypes_df': (mixed_dtypes_df, df_serializer.serialize(mixed_dtypes_df),
-                                df_serializer.can_convert_to_records_without_objects(mixed_dtypes_df, 'symA'))
+                                df_serializer.can_convert_to_records_without_objects(mixed_dtypes_df, 'symA')),
+            'index_tz_aware': (index_tz_aware, df_serializer.serialize(index_tz_aware),
+                               df_serializer.can_convert_to_records_without_objects(index_tz_aware, 'symA'))
         }
     return _TEST_DATA