From f3d7229c98cde50257394dddb41400c44ea051b0 Mon Sep 17 00:00:00 2001 From: James Blackburn Date: Tue, 18 Aug 2015 18:43:06 +0100 Subject: [PATCH] MDP-448 Pass all arguments through from TopLevelTickStore.read to TickStore.read --- arctic/tickstore/toplevel.py | 5 +++-- tests/unit/tickstore/test_toplevel.py | 21 ++++++++++++++++++++- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/arctic/tickstore/toplevel.py b/arctic/tickstore/toplevel.py index 89c984733cb45..12390f031d639 100644 --- a/arctic/tickstore/toplevel.py +++ b/arctic/tickstore/toplevel.py @@ -98,9 +98,10 @@ def add(self, date_range, library_name): self._collection.update_one({'library_name': library_name}, {'$set': {'start': start, 'end': end}}, upsert=True) - def read(self, symbol, date_range, columns=['BID', 'ASK', 'TRDPRC_1', 'BIDSIZE', 'ASKSIZE', 'TRDVOL_1'], **kwargs): + def read(self, symbol, date_range, columns=['BID', 'ASK', 'TRDPRC_1', 'BIDSIZE', 'ASKSIZE', 'TRDVOL_1'], include_images=False): libraries = self._get_libraries(date_range) - dfs = [l.library.read(symbol, l.date_range.intersection(date_range), columns) for l in libraries] + dfs = [l.library.read(symbol, l.date_range.intersection(date_range), columns, + include_images=include_images) for l in libraries] return pd.concat(dfs) def write(self, symbol, data): diff --git a/tests/unit/tickstore/test_toplevel.py b/tests/unit/tickstore/test_toplevel.py index 946f4ecab76ac..b59a71ad24014 100644 --- a/tests/unit/tickstore/test_toplevel.py +++ b/tests/unit/tickstore/test_toplevel.py @@ -1,4 +1,4 @@ -from mock import Mock, patch, MagicMock, create_autospec, sentinel +from mock import Mock, patch, MagicMock, create_autospec, sentinel, call import pytest from datetime import datetime as dt import pandas as pd @@ -10,6 +10,7 @@ from arctic.exceptions import OverlappingDataException from arctic.tickstore.toplevel import TopLevelTickStore, TickStoreLibrary from dateutil.rrule import rrule, DAILY +from arctic.tickstore.tickstore import TickStore def test_raise_exception_if_daterange_is_not_provided(): @@ -145,3 +146,21 @@ def test_write_pandas_data_to_right_libraries(): TopLevelTickStore.write(self, 'blah', sentinel.data) mock_lib1.write.assert_called_once_with('blah', slice1) mock_lib2.write.assert_called_once_with('blah', slice2) + + +def test_read(): + self = create_autospec(TopLevelTickStore) + tsl = TickStoreLibrary(create_autospec(TickStore), create_autospec(DateRange)) + self._get_libraries.return_value = [tsl, tsl] + dr = create_autospec(DateRange) + with patch('pandas.concat') as concat: + res = TopLevelTickStore.read(self, sentinel.symbol, dr, + columns=sentinel.include_columns, + include_images=sentinel.include_images) + assert concat.call_args_list == [call([tsl.library.read.return_value, + tsl.library.read.return_value])] + assert res == concat.return_value + assert tsl.library.read.call_args_list == [call(sentinel.symbol, tsl.date_range.intersection.return_value, + sentinel.include_columns, include_images=sentinel.include_images), + call(sentinel.symbol, tsl.date_range.intersection.return_value, + sentinel.include_columns, include_images=sentinel.include_images)]