test_ts_read.py

# -*- coding: utf-8 -*-
from datetime import datetime as dt

import numpy as np
import pandas as pd
import pytest
import six
from mock import patch, call, Mock
from numpy.testing.utils import assert_array_equal
from pandas import DatetimeIndex
from pandas.util.testing import assert_frame_equal
from pymongo import ReadPreference

from arctic._util import mongo_count
from arctic.date import DateRange, mktz, CLOSED_CLOSED, CLOSED_OPEN, OPEN_CLOSED, OPEN_OPEN
from arctic.exceptions import NoDataFoundException


def test_read(tickstore_lib):
    data = [{'ASK': 1545.25,
                  'ASKSIZE': 1002.0,
                  'BID': 1545.0,
                  'BIDSIZE': 55.0,
                  'CUMVOL': 2187387.0,
                  'DELETED_TIME': 0,
                  'INSTRTYPE': 'FUT',
                  'PRICE': 1545.0,
                  'SIZE': 1.0,
                  'TICK_STATUS': 0,
                  'TRADEHIGH': 1561.75,
                  'TRADELOW': 1537.25,
                  'index': 1185076787070},
                 {'CUMVOL': 354.0,
                  'DELETED_TIME': 0,
                  'PRICE': 1543.75,
                  'SIZE': 354.0,
                  'TRADEHIGH': 1543.75,
                  'TRADELOW': 1543.75,
                  'index': 1185141600600}]
    tickstore_lib.write('FEED::SYMBOL', data)

    df = tickstore_lib.read('FEED::SYMBOL', columns=['BID', 'ASK', 'PRICE'])

    assert_array_equal(df['ASK'].values, np.array([1545.25, np.nan]))
    assert_array_equal(df['BID'].values, np.array([1545, np.nan]))
    assert_array_equal(df['PRICE'].values, np.array([1545, 1543.75]))
    assert_array_equal(df.index.values.astype('object'), np.array([1185076787070000000, 1185141600600000000]))
    assert tickstore_lib._collection.find_one()['c'] == 2
    assert df.index.tzinfo == mktz()


def test_read_data_is_modifiable(tickstore_lib):
    data = [{'ASK': 1545.25,
                  'ASKSIZE': 1002.0,
                  'BID': 1545.0,
                  'BIDSIZE': 55.0,
                  'CUMVOL': 2187387.0,
                  'DELETED_TIME': 0,
                  'INSTRTYPE': 'FUT',
                  'PRICE': 1545.0,
                  'SIZE': 1.0,
                  'TICK_STATUS': 0,
                  'TRADEHIGH': 1561.75,
                  'TRADELOW': 1537.25,
                  'index': 1185076787070},
                 {'CUMVOL': 354.0,
                  'DELETED_TIME': 0,
                  'PRICE': 1543.75,
                  'SIZE': 354.0,
                  'TRADEHIGH': 1543.75,
                  'TRADELOW': 1543.75,
                  'index': 1185141600600}]
    tickstore_lib.write('FEED::SYMBOL', data)

    df = tickstore_lib.read('FEED::SYMBOL', columns=['BID', 'ASK', 'PRICE'])

    df[['BID', 'ASK', 'PRICE']] = 7
    assert np.all(df[['BID', 'ASK', 'PRICE']].values == np.array([[7, 7, 7], [7, 7, 7]]))


def test_read_allow_secondary(tickstore_lib):
    data = [{'ASK': 1545.25,
                  'ASKSIZE': 1002.0,
                  'BID': 1545.0,
                  'BIDSIZE': 55.0,
                  'CUMVOL': 2187387.0,
                  'DELETED_TIME': 0,
                  'INSTRTYPE': 'FUT',
                  'PRICE': 1545.0,
                  'SIZE': 1.0,
                  'TICK_STATUS': 0,
                  'TRADEHIGH': 1561.75,
                  'TRADELOW': 1537.25,
                  'index': 1185076787070},
                 {'CUMVOL': 354.0,
                  'DELETED_TIME': 0,
                  'PRICE': 1543.75,
                  'SIZE': 354.0,
                  'TRADEHIGH': 1543.75,
                  'TRADELOW': 1543.75,
                  'index': 1185141600600}]
    tickstore_lib.write('FEED::SYMBOL', data)

    with patch('pymongo.collection.Collection.find', side_effect=tickstore_lib._collection.find) as find:
        with patch('pymongo.collection.Collection.with_options', side_effect=tickstore_lib._collection.with_options) as with_options:
            with patch.object(tickstore_lib, '_read_preference', side_effect=tickstore_lib._read_preference) as read_pref:
                df = tickstore_lib.read('FEED::SYMBOL', columns=['BID', 'ASK', 'PRICE'], allow_secondary=True)
    assert read_pref.call_args_list == [call(True)]
    assert with_options.call_args_list == [call(read_preference=ReadPreference.NEAREST)]
    assert find.call_args_list == [call({'sy': 'FEED::SYMBOL'}, sort=[('s', 1)], projection={'s': 1, '_id': 0}),
                                   call({'sy': 'FEED::SYMBOL', 's': {'$lte': dt(2007, 8, 21, 3, 59, 47, 70000)}}, 
                                        projection={'sy': 1, 'cs.PRICE': 1, 'i': 1, 'cs.BID': 1, 's': 1, 'im': 1, 'v': 1, 'cs.ASK': 1})]

    assert_array_equal(df['ASK'].values, np.array([1545.25, np.nan]))
    assert tickstore_lib._collection.find_one()['c'] == 2


def test_read_symbol_as_column(tickstore_lib):
    data = [{'ASK': 1545.25,
             'index': 1185076787070},
            {'CUMVOL': 354.0,
             'index': 1185141600600}]
    tickstore_lib.write('FEED::SYMBOL', data)

    df = tickstore_lib.read('FEED::SYMBOL', columns=['SYMBOL', 'CUMVOL'])
    assert all(df['SYMBOL'].values == ['FEED::SYMBOL'])


def test_read_multiple_symbols(tickstore_lib):
    data1 = [{'ASK': 1545.25,
                  'ASKSIZE': 1002.0,
                  'BID': 1545.0,
                  'BIDSIZE': 55.0,
                  'CUMVOL': 2187387.0,
                  'DELETED_TIME': 0,
                  'INSTRTYPE': 'FUT',
                  'PRICE': 1545.0,
                  'SIZE': 1.0,
                  'TICK_STATUS': 0,
                  'TRADEHIGH': 1561.75,
                  'TRADELOW': 1537.25,
                  'index': 1185076787070}, ]
    data2 = [{'CUMVOL': 354.0,
                  'DELETED_TIME': 0,
                  'PRICE': 1543.75,
                  'SIZE': 354.0,
                  'TRADEHIGH': 1543.75,
                  'TRADELOW': 1543.75,
                  'index': 1185141600600}]

    tickstore_lib.write('BAR', data2)
    tickstore_lib.write('FOO', data1)

    df = tickstore_lib.read(['FOO', 'BAR'], columns=['BID', 'ASK', 'PRICE'])

    assert all(df['SYMBOL'].values == ['FOO', 'BAR'])
    assert_array_equal(df['ASK'].values, np.array([1545.25, np.nan]))
    assert_array_equal(df['BID'].values, np.array([1545, np.nan]))
    assert_array_equal(df['PRICE'].values, np.array([1545, 1543.75]))
    assert_array_equal(df.index.values.astype('object'), np.array([1185076787070000000, 1185141600600000000]))
    assert tickstore_lib._collection.find_one()['c'] == 1


@pytest.mark.parametrize('chunk_size', [1, 100])
def test_read_all_cols_all_dtypes(tickstore_lib, chunk_size):
    data = [{'f': 0.1,
            'of': 0.2,
            's': 's',
            'os': 'os',
            'l': 1,
            'ol': 2,
            'index': dt(1970, 1, 1, tzinfo=mktz('UTC')),
            },
            {'f': 0.3,
            'nf': 0.4,
            's': 't',
            'ns': 'ns',
            'l': 3,
            'nl': 4,
            'index': dt(1970, 1, 1, 0, 0, 1, tzinfo=mktz('UTC')),
            },
            ]
    tickstore_lib._chunk_size = chunk_size
    tickstore_lib.write('sym', data)
    df = tickstore_lib.read('sym', columns=None)

    assert df.index.tzinfo == mktz()

    # The below is probably more trouble than it's worth, but we *should*
    # be able to roundtrip data and get the same answer...

    # Ints become floats
    data[0]['l'] = float(data[0]['l'])
    # Treat missing strings as None
    data[0]['ns'] = None
    data[1]['os'] = None
    index = DatetimeIndex([dt(1970, 1, 1, tzinfo=mktz('UTC')),
                         dt(1970, 1, 1, 0, 0, 1, tzinfo=mktz('UTC'))],
                        )
    df.index = df.index.tz_convert(mktz('UTC'))
    expected = pd.DataFrame(data, index=index)
    expected = expected[df.columns]
    assert_frame_equal(expected, df, check_names=False)


DUMMY_DATA = [
              {'a': 1.,
               'b': 2.,
               'index': dt(2013, 1, 1, tzinfo=mktz('Europe/London'))
               },
              {'b': 3.,
               'c': 4.,
               'index': dt(2013, 1, 2, tzinfo=mktz('Europe/London'))
               },
              {'b': 5.,
               'c': 6.,
               'index': dt(2013, 1, 3, tzinfo=mktz('Europe/London'))
               },
              {'b': 7.,
               'c': 8.,
               'index': dt(2013, 1, 4, tzinfo=mktz('Europe/London'))
               },
              {'b': 9.,
               'c': 10.,
               'index': dt(2013, 1, 5, tzinfo=mktz('Europe/London'))
               },
              ]


def test_date_range(tickstore_lib):
    tickstore_lib.write('SYM', DUMMY_DATA)
    df = tickstore_lib.read('SYM', date_range=DateRange(20130101, 20130103), columns=None)
    assert_array_equal(df['a'].values, np.array([1, np.nan, np.nan]))
    assert_array_equal(df['b'].values, np.array([2., 3., 5.]))
    assert_array_equal(df['c'].values, np.array([np.nan, 4., 6.]))

    tickstore_lib.delete('SYM')

    # Chunk every 3 symbols and lets have some fun
    tickstore_lib._chunk_size = 3
    tickstore_lib.write('SYM', DUMMY_DATA)

    with patch('pymongo.collection.Collection.find', side_effect=tickstore_lib._collection.find) as f:
        df = tickstore_lib.read('SYM', date_range=DateRange(20130101, 20130103), columns=None)
        assert_array_equal(df['b'].values, np.array([2., 3., 5.]))
        assert mongo_count(tickstore_lib._collection, filter=f.call_args_list[-1][0][0]) == 1
        df = tickstore_lib.read('SYM', date_range=DateRange(20130102, 20130103), columns=None)
        assert_array_equal(df['b'].values, np.array([3., 5.]))
        assert mongo_count(tickstore_lib._collection, filter=f.call_args_list[-1][0][0]) == 1
        df = tickstore_lib.read('SYM', date_range=DateRange(20130103, 20130103), columns=None)
        assert_array_equal(df['b'].values, np.array([5.]))
        assert mongo_count(tickstore_lib._collection, filter=f.call_args_list[-1][0][0]) == 1

        df = tickstore_lib.read('SYM', date_range=DateRange(20130102, 20130104), columns=None)
        assert_array_equal(df['b'].values, np.array([3., 5., 7.]))
        assert mongo_count(tickstore_lib._collection, filter=f.call_args_list[-1][0][0]) == 2
        df = tickstore_lib.read('SYM', date_range=DateRange(20130102, 20130105), columns=None)
        assert_array_equal(df['b'].values, np.array([3., 5., 7., 9.]))
        assert mongo_count(tickstore_lib._collection, filter=f.call_args_list[-1][0][0]) == 2

        df = tickstore_lib.read('SYM', date_range=DateRange(20130103, 20130104), columns=None)
        assert_array_equal(df['b'].values, np.array([5., 7.]))
        assert mongo_count(tickstore_lib._collection, filter=f.call_args_list[-1][0][0]) == 2
        df = tickstore_lib.read('SYM', date_range=DateRange(20130103, 20130105), columns=None)
        assert_array_equal(df['b'].values, np.array([5., 7., 9.]))
        assert mongo_count(tickstore_lib._collection, filter=f.call_args_list[-1][0][0]) == 2

        df = tickstore_lib.read('SYM', date_range=DateRange(20130104, 20130105), columns=None)
        assert_array_equal(df['b'].values, np.array([7., 9.]))
        assert mongo_count(tickstore_lib._collection, filter=f.call_args_list[-1][0][0]) == 1

        # Test the different open-closed behaviours
        df = tickstore_lib.read('SYM', date_range=DateRange(20130104, 20130105, CLOSED_CLOSED), columns=None)
        assert_array_equal(df['b'].values, np.array([7., 9.]))
        df = tickstore_lib.read('SYM', date_range=DateRange(20130104, 20130105, CLOSED_OPEN), columns=None)
        assert_array_equal(df['b'].values, np.array([7.]))
        df = tickstore_lib.read('SYM', date_range=DateRange(20130104, 20130105, OPEN_CLOSED), columns=None)
        assert_array_equal(df['b'].values, np.array([9.]))
        df = tickstore_lib.read('SYM', date_range=DateRange(20130104, 20130105, OPEN_OPEN), columns=None)
        assert_array_equal(df['b'].values, np.array([]))


def test_date_range_end_not_in_range(tickstore_lib):
    DUMMY_DATA = [
                  {'a': 1.,
                   'b': 2.,
                   'index': dt(2013, 1, 1, tzinfo=mktz('Europe/London'))
                   },
                  {'b': 3.,
                   'c': 4.,
                   'index': dt(2013, 1, 2, 10, 1, tzinfo=mktz('Europe/London'))
                   },
                  ]

    tickstore_lib._chunk_size = 1
    tickstore_lib.write('SYM', DUMMY_DATA)
    with patch.object(tickstore_lib._collection, 'find', side_effect=tickstore_lib._collection.find) as f:
        df = tickstore_lib.read('SYM', date_range=DateRange(20130101, dt(2013, 1, 2, 9, 0)), columns=None)
        assert_array_equal(df['b'].values, np.array([2.]))
        assert mongo_count(tickstore_lib._collection, filter=f.call_args_list[-1][0][0]) == 1


@pytest.mark.parametrize('tz_name', ['UTC',
                                     'Europe/London',  # Sometimes ahead of UTC
                                     'America/New_York',  # Behind UTC
                                      ])
def test_date_range_default_timezone(tickstore_lib, tz_name):
    """
    We assume naive datetimes are user-local
    """
    DUMMY_DATA = [
                  {'a': 1.,
                   'b': 2.,
                   'index': dt(2013, 1, 1, tzinfo=mktz(tz_name))
                   },
                  # Half-way through the year
                  {'b': 3.,
                   'c': 4.,
                   'index': dt(2013, 7, 1, tzinfo=mktz(tz_name))
                   },
                  ]

    with patch('tzlocal.get_localzone', return_value=Mock(zone=tz_name)):
        tickstore_lib._chunk_size = 1
        tickstore_lib.write('SYM', DUMMY_DATA)
        df = tickstore_lib.read('SYM', date_range=DateRange(20130101, 20130701), columns=None)

        assert df.index.tzinfo == mktz()

        assert len(df) == 2
        assert df.index[1] == dt(2013, 7, 1, tzinfo=mktz(tz_name))
        df = tickstore_lib.read('SYM', date_range=DateRange(20130101, 20130101), columns=None)
        assert len(df) == 1
        assert df.index.tzinfo == mktz()

        df = tickstore_lib.read('SYM', date_range=DateRange(20130701, 20130701), columns=None)
        assert len(df) == 1
        assert df.index.tzinfo == mktz()


def test_date_range_no_bounds(tickstore_lib):
    DUMMY_DATA = [
                  {'a': 1.,
                   'b': 2.,
                   'index': dt(2013, 1, 1, tzinfo=mktz('Europe/London'))
                   },
                  {'a': 3.,
                   'b': 4.,
                   'index': dt(2013, 1, 30, tzinfo=mktz('Europe/London'))
                   },
                  {'b': 5.,
                   'c': 6.,
                   'index': dt(2013, 2, 2, 10, 1, tzinfo=mktz('Europe/London'))
                   },
                  ]

    tickstore_lib._chunk_size = 1
    tickstore_lib.write('SYM', DUMMY_DATA)

    # 1) No start, no end
    df = tickstore_lib.read('SYM', columns=None)
    assert_array_equal(df['b'].values, np.array([2., 4.]))
    # 1.2) Start before the real start
    df = tickstore_lib.read('SYM', date_range=DateRange(20121231), columns=None)
    assert_array_equal(df['b'].values, np.array([2., 4.]))
    # 2.1) Only go one month out
    df = tickstore_lib.read('SYM', date_range=DateRange(20130101), columns=None)
    assert_array_equal(df['b'].values, np.array([2., 4.]))
    # 2.2) Only go one month out
    df = tickstore_lib.read('SYM', date_range=DateRange(20130102), columns=None)
    assert_array_equal(df['b'].values, np.array([4.]))
    # 3) No start
    df = tickstore_lib.read('SYM', date_range=DateRange(end=20130102), columns=None)
    assert_array_equal(df['b'].values, np.array([2.]))
    # 4) Outside bounds
    df = tickstore_lib.read('SYM', date_range=DateRange(end=20131212), columns=None)
    assert_array_equal(df['b'].values, np.array([2., 4., 5.]))


def test_date_range_BST(tickstore_lib):
    DUMMY_DATA = [
                  {'a': 1.,
                   'b': 2.,
                   'index': dt(2013, 6, 1, 12, 00, tzinfo=mktz('Europe/London'))
                   },
                  {'a': 3.,
                   'b': 4.,
                   'index': dt(2013, 6, 1, 13, 00, tzinfo=mktz('Europe/London'))
                   },
                  ]
    tickstore_lib._chunk_size = 1
    tickstore_lib.write('SYM', DUMMY_DATA)

    df = tickstore_lib.read('SYM', columns=None)
    assert_array_equal(df['b'].values, np.array([2., 4.]))

#     df = tickstore_lib.read('SYM', columns=None, date_range=DateRange(dt(2013, 6, 1, 12),
#                                                                       dt(2013, 6, 1, 13)))
#     assert_array_equal(df['b'].values, np.array([2., 4.]))
    df = tickstore_lib.read('SYM', columns=None, date_range=DateRange(dt(2013, 6, 1, 12, tzinfo=mktz('Europe/London')),
                                                                            dt(2013, 6, 1, 13, tzinfo=mktz('Europe/London'))))
    assert_array_equal(df['b'].values, np.array([2., 4.]))

    df = tickstore_lib.read('SYM', columns=None, date_range=DateRange(dt(2013, 6, 1, 12, tzinfo=mktz('UTC')),
                                                                            dt(2013, 6, 1, 13, tzinfo=mktz('UTC'))))
    assert_array_equal(df['b'].values, np.array([4., ]))


def test_read_no_data(tickstore_lib):
    with pytest.raises(NoDataFoundException):
        tickstore_lib.read('missing_sym', DateRange(20131212, 20131212))


def test_write_no_tz(tickstore_lib):
    DUMMY_DATA = [
                  {'a': 1.,
                   'b': 2.,
                   'index': dt(2013, 6, 1, 12, 00)
                   }]
    with pytest.raises(ValueError):
        tickstore_lib.write('SYM', DUMMY_DATA)


def test_read_out_of_order(tickstore_lib):
    data = [{'A': 120, 'D': 1}, {'A': 122, 'B': 2.0}, {'A': 3, 'B': 3.0, 'D': 1}]
    tick_index = [dt(2013, 6, 1, 12, 00, tzinfo=mktz('UTC')),
                  dt(2013, 6, 1, 11, 00, tzinfo=mktz('UTC')),  # Out-of-order
                  dt(2013, 6, 1, 13, 00, tzinfo=mktz('UTC'))]
    data = pd.DataFrame(data, index=tick_index)

    tickstore_lib._chunk_size = 3
    tickstore_lib.write('SYM', data)
    tickstore_lib.read('SYM', columns=None)
    assert len(tickstore_lib.read('SYM', columns=None, date_range=DateRange(dt(2013, 6, 1, tzinfo=mktz('UTC')), dt(2013, 6, 2, tzinfo=mktz('UTC'))))) == 3
    assert len(tickstore_lib.read('SYM', columns=None, date_range=DateRange(dt(2013, 6, 1, tzinfo=mktz('UTC')), dt(2013, 6, 1, 12, tzinfo=mktz('UTC'))))) == 2


def test_read_chunk_boundaries(tickstore_lib):
    SYM1_DATA = [
                  {'a': 1.,
                   'b': 2.,
                   'index': dt(2013, 6, 1, 12, 00, tzinfo=mktz('UTC'))
                   },
                   {'a': 3.,
                   'b': 4.,
                   'index': dt(2013, 6, 1, 13, 00, tzinfo=mktz('UTC'))
                   },
                 # Chunk boundary here
                   {'a': 5.,
                   'b': 6.,
                   'index': dt(2013, 6, 1, 14, 00, tzinfo=mktz('UTC'))
                   }
                  ]
    SYM2_DATA = [
                  {'a': 7.,
                   'b': 8.,
                   'index': dt(2013, 6, 1, 12, 30, tzinfo=mktz('UTC'))
                   },
                   {'a': 9.,
                   'b': 10.,
                   'index': dt(2013, 6, 1, 13, 30, tzinfo=mktz('UTC'))
                   },
                 # Chunk boundary here
                   {'a': 11.,
                   'b': 12.,
                   'index': dt(2013, 6, 1, 14, 30, tzinfo=mktz('UTC'))
                   }
                  ]
    tickstore_lib._chunk_size = 2
    tickstore_lib.write('SYM1', SYM1_DATA)
    tickstore_lib.write('SYM2', SYM2_DATA)

    assert len(tickstore_lib.read('SYM1', columns=None, date_range=DateRange(dt(2013, 6, 1, 12, 45, tzinfo=mktz('UTC')), dt(2013, 6, 1, 15, 00, tzinfo=mktz('UTC'))))) == 2
    assert len(tickstore_lib.read('SYM2', columns=None, date_range=DateRange(dt(2013, 6, 1, 12, 45, tzinfo=mktz('UTC')), dt(2013, 6, 1, 15, 00, tzinfo=mktz('UTC'))))) == 2

    assert len(tickstore_lib.read(['SYM1', 'SYM2'], columns=None, date_range=DateRange(dt(2013, 6, 1, 12, 45, tzinfo=mktz('UTC')), dt(2013, 6, 1, 15, 00, tzinfo=mktz('UTC'))))) == 4


def test_read_spanning_chunks(tickstore_lib):
    SYM1_DATA = [
              {'a': 1.,
               'b': 2.,
               'index': dt(2013, 6, 1, 11, 00, tzinfo=mktz('UTC'))
               },
               {'a': 3.,
               'b': 4.,
               'index': dt(2013, 6, 1, 12, 00, tzinfo=mktz('UTC'))
               },
             # Chunk boundary here
               {'a': 5.,
               'b': 6.,
               'index': dt(2013, 6, 1, 14, 00, tzinfo=mktz('UTC'))
               }
              ]
    SYM2_DATA = [
                  {'a': 7.,
                   'b': 8.,
                   'index': dt(2013, 6, 1, 12, 30, tzinfo=mktz('UTC'))
                   },
                   {'a': 9.,
                   'b': 10.,
                   'index': dt(2013, 6, 1, 13, 30, tzinfo=mktz('UTC'))
                   },
                 # Chunk boundary here
                   {'a': 11.,
                   'b': 12.,
                   'index': dt(2013, 6, 1, 14, 30, tzinfo=mktz('UTC'))
                   }
                  ]
    tickstore_lib._chunk_size = 2
    tickstore_lib.write('SYM1', SYM1_DATA)
    tickstore_lib.write('SYM2', SYM2_DATA)

    # Even though the latest chunk that's the closest to the start point for SYM1 starts at 11:00, it ends before the start point,
    # so we want to ignore it and start from SYM2 (12:30) instead.
    assert tickstore_lib._mongo_date_range_query(
                                                 ['SYM1', 'SYM2'],
                                                 date_range=DateRange(dt(2013, 6, 1, 12, 45, tzinfo=mktz('UTC')),
                                                                      dt(2013, 6, 1, 15, 00, tzinfo=mktz('UTC')))) == \
        {'s': {'$gte': dt(2013, 6, 1, 12, 30, tzinfo=mktz('UTC')), '$lte': dt(2013, 6, 1, 15, 0, tzinfo=mktz('UTC'))}}


def test_read_inside_range(tickstore_lib):
    SYM1_DATA = [
              {'a': 1.,
               'b': 2.,
               'index': dt(2013, 6, 1, 0, 00, tzinfo=mktz('UTC'))
               },
               {'a': 3.,
               'b': 4.,
               'index': dt(2013, 6, 1, 1, 00, tzinfo=mktz('UTC'))
               },
             # Chunk boundary here
               {'a': 5.,
               'b': 6.,
               'index': dt(2013, 6, 1, 14, 00, tzinfo=mktz('UTC'))
               }
              ]
    SYM2_DATA = [
                  {'a': 7.,
                   'b': 8.,
                   'index': dt(2013, 6, 1, 12, 30, tzinfo=mktz('UTC'))
                   },
                   {'a': 9.,
                   'b': 10.,
                   'index': dt(2013, 6, 1, 13, 30, tzinfo=mktz('UTC'))
                   },
                 # Chunk boundary here
                   {'a': 11.,
                   'b': 12.,
                   'index': dt(2013, 6, 1, 14, 30, tzinfo=mktz('UTC'))
                   }
                  ]
    tickstore_lib._chunk_size = 2
    tickstore_lib.write('SYM1', SYM1_DATA)
    tickstore_lib.write('SYM2', SYM2_DATA)

    # If there are no chunks spanning the range, we still cap the start range so that we don't
    # fetch SYM1's 0am--1am chunk
    assert tickstore_lib._mongo_date_range_query(
                                                 ['SYM1', 'SYM2'],
                                                 date_range=DateRange(dt(2013, 6, 1, 10, 0, tzinfo=mktz('UTC')),
                                                                      dt(2013, 6, 1, 15, 0, tzinfo=mktz('UTC')))) == \
        {'s': {'$gte': dt(2013, 6, 1, 10, 0, tzinfo=mktz('UTC')), '$lte': dt(2013, 6, 1, 15, 0, tzinfo=mktz('UTC'))}}

def test_read_longs(tickstore_lib):
    DUMMY_DATA = [
                  {'a': 1,
                   'index': dt(2013, 6, 1, 12, 00, tzinfo=mktz('Europe/London'))
                   },
                  {
                   'b': 4,
                   'index': dt(2013, 6, 1, 13, 00, tzinfo=mktz('Europe/London'))
                   },
                  ]
    tickstore_lib._chunk_size = 3
    tickstore_lib.write('SYM', DUMMY_DATA)
    tickstore_lib.read('SYM', columns=None)
    read = tickstore_lib.read('SYM', columns=None, date_range=DateRange(dt(2013, 6, 1), dt(2013, 6, 2)))
    assert read['a'][0] == 1
    assert np.isnan(read['b'][0])


def test_read_with_image(tickstore_lib):
    DUMMY_DATA = [
              {'a': 1.,
               'index': dt(2013, 1, 1, 11, 00, tzinfo=mktz('Europe/London'))
               },
              {
               'b': 4.,
               'index': dt(2013, 1, 1, 12, 00, tzinfo=mktz('Europe/London'))
               },
              ]
    # Add an image
    tickstore_lib.write('SYM', DUMMY_DATA)
    tickstore_lib._collection.update_one({},
                                     {'$set':
                                      {'im': {'i':
                                              {'a': 37.,
                                               'c': 2.,
                                               },
                                              't': dt(2013, 1, 1, 10, tzinfo=mktz('Europe/London'))
                                              }
                                       }
                                      }
                                     )

    dr = DateRange(dt(2013, 1, 1), dt(2013, 1, 2))
    # tickstore_lib.read('SYM', columns=None)
    df = tickstore_lib.read('SYM', columns=None, date_range=dr)
    assert df['a'][0] == 1

    # Read with the image as well - all columns
    df = tickstore_lib.read('SYM', columns=None, date_range=dr, include_images=True)
    assert set(df.columns) == set(('a', 'b', 'c'))
    assert_array_equal(df['a'].values, np.array([37, 1, np.nan]))
    assert_array_equal(df['b'].values, np.array([np.nan, np.nan, 4]))
    assert_array_equal(df['c'].values, np.array([2, np.nan, np.nan]))
    assert df.index[0] == dt(2013, 1, 1, 10, tzinfo=mktz('Europe/London'))
    assert df.index[1] == dt(2013, 1, 1, 11, tzinfo=mktz('Europe/London'))
    assert df.index[2] == dt(2013, 1, 1, 12, tzinfo=mktz('Europe/London'))

    # Read just columns from the updates
    df = tickstore_lib.read('SYM', columns=('a', 'b'), date_range=dr, include_images=True)
    assert set(df.columns) == set(('a', 'b'))
    assert_array_equal(df['a'].values, np.array([37, 1, np.nan]))
    assert_array_equal(df['b'].values, np.array([np.nan, np.nan, 4]))
    assert df.index[0] == dt(2013, 1, 1, 10, tzinfo=mktz('Europe/London'))
    assert df.index[1] == dt(2013, 1, 1, 11, tzinfo=mktz('Europe/London'))
    assert df.index[2] == dt(2013, 1, 1, 12, tzinfo=mktz('Europe/London'))

    # Read one column from the updates
    df = tickstore_lib.read('SYM', columns=('a',), date_range=dr, include_images=True)
    assert set(df.columns) == set(('a',))
    assert_array_equal(df['a'].values, np.array([37, 1]))
    assert df.index[0] == dt(2013, 1, 1, 10, tzinfo=mktz('Europe/London'))
    assert df.index[1] == dt(2013, 1, 1, 11, tzinfo=mktz('Europe/London'))

    # Read just the image column
    df = tickstore_lib.read('SYM', columns=['c'], date_range=dr, include_images=True)
    assert set(df.columns) == set(['c'])
    assert_array_equal(df['c'].values, np.array([2]))
    assert df.index[0] == dt(2013, 1, 1, 10, tzinfo=mktz('Europe/London'))


def test_read_with_metadata(tickstore_lib):
    metadata = {'metadata': 'important data'}
    tickstore_lib.write('test', [{'index': dt(2013, 6, 1, 13, 00, tzinfo=mktz('Europe/London')), 'price': 100.50, 'ticker': 'QQQ'}], metadata=metadata)
    m = tickstore_lib.read_metadata('test')
    assert(metadata == m)


def test_read_strings(tickstore_lib):
    df = pd.DataFrame(data={'data': ['A', 'B', 'C']},
                      index=pd.Index(data=[dt(2016, 1, 1, 00, tzinfo=mktz('UTC')),
                                           dt(2016, 1, 2, 00, tzinfo=mktz('UTC')),
                                           dt(2016, 1, 3, 00, tzinfo=mktz('UTC'))], name='date'))
    tickstore_lib.write('test', df)
    read_df = tickstore_lib.read('test')
    assert(all(read_df['data'].values == df['data'].values))


def test_read_utf8_strings(tickstore_lib):
    data = ['一', '二', '三'] # Chinese character [one, two , three]
    if six.PY2:
      utf8_data = data
      unicode_data = [s.decode('utf8') for s in data]
    else:
      utf8_data = [s.encode('utf8') for s in data]
      unicode_data = data
    df = pd.DataFrame(data={'data': utf8_data},
                      index=pd.Index(data=[dt(2016, 1, 1, 00, tzinfo=mktz('UTC')),
                                           dt(2016, 1, 2, 00, tzinfo=mktz('UTC')),
                                           dt(2016, 1, 3, 00, tzinfo=mktz('UTC'))], name='date'))
    tickstore_lib.write('test', df)
    read_df = tickstore_lib.read('test')
    assert(all(read_df['data'].values == np.array(unicode_data)))


def test_read_unicode_strings(tickstore_lib):
    df = pd.DataFrame(data={'data': [u'一', u'二', u'三']}, # Chinese character [one, two , three]
                      index=pd.Index(data=[dt(2016, 1, 1, 00, tzinfo=mktz('UTC')),
                                           dt(2016, 1, 2, 00, tzinfo=mktz('UTC')),
                                           dt(2016, 1, 3, 00, tzinfo=mktz('UTC'))], name='date'))
    tickstore_lib.write('test', df)
    read_df = tickstore_lib.read('test')
    assert(all(read_df['data'].values == df['data'].values))


def test_objects_fail(tickstore_lib):
    class Fake(object):
        def __init__(self, val):
            self.val = val

        def fake(self):
            return self.val

    df = pd.DataFrame(data={'data': [Fake(1), Fake(2)]},
                      index=pd.Index(data=[dt(2016, 1, 1, 00, tzinfo=mktz('UTC')),
                                           dt(2016, 1, 2, 00, tzinfo=mktz('UTC'))], name='date'))

    with pytest.raises(Exception) as e:
        tickstore_lib.write('test', df)
    assert('Casting object column to string failed' in str(e.value))