From 5b37af6e04506402e391892a5ed41ae0bdcf40c7 Mon Sep 17 00:00:00 2001 From: Maya Tydykov Date: Mon, 22 Feb 2016 14:09:39 -0500 Subject: [PATCH] TST: enhance test; add more common data. STY: fixes for PEP8. --- tests/pipeline/test_buyback_auth.py | 37 ++++---- tests/pipeline/test_earnings.py | 66 +++++++------- tests/pipeline/test_events.py | 130 ++++++++++++++++++++-------- 3 files changed, 143 insertions(+), 90 deletions(-) diff --git a/tests/pipeline/test_buyback_auth.py b/tests/pipeline/test_buyback_auth.py index 10908f87..6d4ee10b 100644 --- a/tests/pipeline/test_buyback_auth.py +++ b/tests/pipeline/test_buyback_auth.py @@ -8,13 +8,10 @@ import blaze as bz from blaze.compute.core import swap_resources_into_scope from contextlib2 import ExitStack from nose_parameterized import parameterized -import numpy as np import pandas as pd -from pandas.util.testing import assert_series_equal from six import iteritems from tests.pipeline.test_events import param_dates, EventLoaderCommonTest -from zipline.pipeline import Pipeline from zipline.pipeline.common import( BUYBACK_ANNOUNCEMENT_FIELD_NAME, CASH_FIELD_NAME, @@ -27,7 +24,6 @@ from zipline.pipeline.common import( TS_FIELD_NAME) from zipline.pipeline.data import (CashBuybackAuthorizations, ShareBuybackAuthorizations) -from zipline.pipeline.engine import SimplePipelineEngine from zipline.pipeline.factors.events import ( BusinessDaysSincePreviousCashBuybackAuth, BusinessDaysSincePreviousShareBuybackAuth @@ -38,11 +34,7 @@ from zipline.pipeline.loaders.blaze import ( BlazeCashBuybackAuthorizationsLoader, BlazeShareBuybackAuthorizationsLoader, ) -from zipline.utils.numpy_utils import make_datetime64D, NaTD from zipline.utils.test_utils import ( - gen_calendars, - make_simple_equity_info, - num_days_in_range, tmp_asset_finder, ) @@ -97,7 +89,7 @@ class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonTest): Test for cash buyback authorizations dataset. """ pipeline_columns = { - ('%s' % PREVIOUS_BUYBACK_CASH): + (PREVIOUS_BUYBACK_CASH): CashBuybackAuthorizations.previous_value.latest, PREVIOUS_BUYBACK_ANNOUNCEMENT: CashBuybackAuthorizations.previous_announcement_date.latest, @@ -124,8 +116,6 @@ class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonTest): zip_with_floats_dates = partial(self.zip_with_floats, dates) num_days_between_dates = partial(self.num_days_between, dates) _expected_previous_cash = pd.DataFrame({ - # TODO if the next knowledge date is 10, why is the range - # until 15? 0: zip_with_floats_dates( ['NaN'] * num_days_between_dates(None, '2014-01-14') + [10] * num_days_between_dates('2014-01-15', '2014-01-19') + @@ -148,14 +138,16 @@ class CashBuybackAuthLoaderTestCase(TestCase, EventLoaderCommonTest): ), 4: zip_with_floats_dates(['NaN'] * len(dates)), }, index=dates) - self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] = self.get_expected_previous( - dates) + self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] = \ + self.get_expected_previous_event_dates(dates) self.cols[PREVIOUS_BUYBACK_CASH] = _expected_previous_cash - + self.cols[DAYS_SINCE_PREV] = self._compute_busday_offsets( + self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] + ) @parameterized.expand(param_dates) def test_compute_cash_buyback_auth(self, dates): - self._test_compute_buyback_auth(dates) + self._test_compute(dates) class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase): @@ -163,9 +155,9 @@ class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase): Test for share buyback authorizations dataset. """ pipeline_columns = { - ('%s' % PREVIOUS_BUYBACK_SHARE_COUNT): + PREVIOUS_BUYBACK_SHARE_COUNT: ShareBuybackAuthorizations.previous_share_count.latest, - ('%s' % PREVIOUS_BUYBACK_ANNOUNCEMENT): + PREVIOUS_BUYBACK_ANNOUNCEMENT: ShareBuybackAuthorizations.previous_announcement_date.latest, DAYS_SINCE_PREV: BusinessDaysSincePreviousShareBuybackAuth(), @@ -179,8 +171,8 @@ class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase): ) cls.cols = {} cls.dataset = {sid: df.drop(CASH_FIELD_NAME, 1) - for sid, df in - enumerate(buyback_authorizations)} + for sid, df in + enumerate(buyback_authorizations)} cls.loader_type = ShareBuybackAuthorizationsLoader @classmethod @@ -217,11 +209,14 @@ class ShareBuybackAuthLoaderTestCase(EventLoaderCommonTest, TestCase): PREVIOUS_BUYBACK_SHARE_COUNT ] = _expected_previous_buyback_share_count self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] = \ - self.get_expected_previous(dates) + self.get_expected_previous_event_dates(dates) + self.cols[DAYS_SINCE_PREV] = self._compute_busday_offsets( + self.cols[PREVIOUS_BUYBACK_ANNOUNCEMENT] + ) @parameterized.expand(param_dates) def test_compute_share_buyback_auth(self, dates): - self._test_compute_buyback_auth(dates) + self._test_compute(dates) class BlazeCashBuybackAuthLoaderTestCase(CashBuybackAuthLoaderTestCase): diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py index f5242267..a4bce18b 100644 --- a/tests/pipeline/test_earnings.py +++ b/tests/pipeline/test_earnings.py @@ -35,35 +35,35 @@ from zipline.utils.test_utils import ( ) earnings_dates = [ - # K1--K2--E1--E2. - pd.DataFrame({ - TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']), - ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-15', - '2014-01-20']) - }), - # K1--K2--E2--E1. - pd.DataFrame({ - TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']), - ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-20', - '2014-01-15']) - }), - # K1--E1--K2--E2. - pd.DataFrame({ - TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-15']), - ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-10', - '2014-01-20']) - }), - # K1 == K2. - pd.DataFrame({ - TS_FIELD_NAME: pd.to_datetime(['2014-01-05'] * 2), - ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-10', - '2014-01-15']) - }), - pd.DataFrame({ - TS_FIELD_NAME: pd.to_datetime([]), - ANNOUNCEMENT_FIELD_NAME: pd.to_datetime([]) - }) - ] + # K1--K2--E1--E2. + pd.DataFrame({ + TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']), + ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-15', + '2014-01-20']) + }), + # K1--K2--E2--E1. + pd.DataFrame({ + TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-10']), + ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-20', + '2014-01-15']) + }), + # K1--E1--K2--E2. + pd.DataFrame({ + TS_FIELD_NAME: pd.to_datetime(['2014-01-05', '2014-01-15']), + ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-10', + '2014-01-20']) + }), + # K1 == K2. + pd.DataFrame({ + TS_FIELD_NAME: pd.to_datetime(['2014-01-05'] * 2), + ANNOUNCEMENT_FIELD_NAME: pd.to_datetime(['2014-01-10', + '2014-01-15']) + }), + pd.DataFrame({ + TS_FIELD_NAME: pd.to_datetime([]), + ANNOUNCEMENT_FIELD_NAME: pd.to_datetime([]) + }) +] class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonTest): @@ -97,11 +97,12 @@ class EarningsCalendarLoaderTestCase(TestCase, EventLoaderCommonTest): def tearDownClass(cls): cls._cleanup_stack.close() - def setup(self, dates): _expected_next_announce = self.get_expected_next_event_dates(dates) - _expected_previous_announce = self.get_expected_previous_event_dates(dates) + _expected_previous_announce = self.get_expected_previous_event_dates( + dates + ) _expected_next_busday_offsets = self._compute_busday_offsets( _expected_next_announce @@ -146,7 +147,8 @@ class BlazeEarningsCalendarLoaderNotInteractiveTestCase( """ @classmethod def setUpClass(cls): - super(BlazeEarningsCalendarLoaderNotInteractiveTestCase, cls).setUpClass() + super(BlazeEarningsCalendarLoaderNotInteractiveTestCase, + cls).setUpClass() cls.loader_type = BlazeEarningsCalendarLoader def loader_args(self, dates): diff --git a/tests/pipeline/test_events.py b/tests/pipeline/test_events.py index 33c328ec..6c94b902 100644 --- a/tests/pipeline/test_events.py +++ b/tests/pipeline/test_events.py @@ -26,7 +26,9 @@ from zipline.pipeline.loaders.events import ( WRONG_COLS_ERROR, ) from zipline.utils.memoize import lazyval -from zipline.utils.numpy_utils import datetime64ns_dtype, NaTD, make_datetime64D +from zipline.utils.numpy_utils import (datetime64ns_dtype, + NaTD, + make_datetime64D) from zipline.utils.test_utils import gen_calendars, num_days_in_range, \ make_simple_equity_info @@ -153,25 +155,42 @@ class EventLoaderTestCase(TestCase): expected, ) - @parameterized.expand([ - # DataFrame without timestamp column and infer_timestamps = True - [pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx}), - False, - DF_NO_TS_NOT_INFER_TS_ERROR.format( - timestamp_column_name=TS_FIELD_NAME, - sid=0 - ) - ], - # DatetimeIndex with infer_timestamps = False - [pd.DatetimeIndex(dtx, name=ANNOUNCEMENT_FIELD_NAME), False, - DTINDEX_NOT_INFER_TS_ERROR.format(sid=0)], - # Series with DatetimeIndex as index and infer_timestamps = False - [pd.Series(dtx, name=ANNOUNCEMENT_FIELD_NAME), False, - SERIES_NO_DTINDEX_ERROR.format(sid=0)], - # Some other data structure that is not expected - [dtx, False, BAD_DATA_FORMAT_ERROR.format(sid=0)], - [dtx, True, BAD_DATA_FORMAT_ERROR.format(sid=0)] - ]) + @parameterized.expand( + [ + # DataFrame without timestamp column and infer_timestamps = True + [ + pd.DataFrame({ANNOUNCEMENT_FIELD_NAME: dtx}), + False, + DF_NO_TS_NOT_INFER_TS_ERROR.format( + timestamp_column_name=TS_FIELD_NAME, + sid=0 + ) + ], + # DatetimeIndex with infer_timestamps = False + [ + pd.DatetimeIndex(dtx, name=ANNOUNCEMENT_FIELD_NAME), + False, + DTINDEX_NOT_INFER_TS_ERROR.format(sid=0) + ], + # Series with DatetimeIndex as index and infer_timestamps = False + [ + pd.Series(dtx, name=ANNOUNCEMENT_FIELD_NAME), + False, + SERIES_NO_DTINDEX_ERROR.format(sid=0) + ], + # Some other data structure that is not expected + [ + dtx, + False, + BAD_DATA_FORMAT_ERROR.format(sid=0) + ], + [ + dtx, + True, + BAD_DATA_FORMAT_ERROR.format(sid=0) + ] + ] + ) def test_bad_conversion_to_df(self, df, infer_timestamps, msg): events_by_sid = {0: df} assert_loader_error(events_by_sid, ValueError, msg, @@ -204,14 +223,6 @@ class BlazeEventLoaderTestCase(TestCase): TestCase.assertTrue(ABSTRACT_METHODS_ERROR in context.exception) - - - - - -########################## - - # Must be a list - can't use generator since this needs to be used more than # once. param_dates = list(gen_calendars( @@ -237,15 +248,12 @@ class EventLoaderCommonTest(object): def zip_with_floats(self, dates, flts): return pd.Series(flts, index=dates).astype('float') - def num_days_between(self, dates, start_date, end_date): return num_days_in_range(dates, start_date, end_date) - def zip_with_dates(self, index_dates, dts): return pd.Series(pd.to_datetime(dts), index=index_dates) - def loader_args(self, dates): """Construct the base object to pass to the loader. @@ -268,7 +276,59 @@ class EventLoaderCommonTest(object): loader = self.loader_type(*self.loader_args(dates)) return SimplePipelineEngine(lambda _: loader, dates, self.finder) - def get_expected_previous(self, dates): + def get_expected_next_event_dates(self, dates): + num_days_between_for_dates = partial(self.num_days_between, dates) + zip_with_dates_for_dates = partial(self.zip_with_dates, dates) + return pd.DataFrame({ + 0: zip_with_dates_for_dates( + ['NaT'] * + num_days_between_for_dates(None, '2014-01-04') + + ['2014-01-15'] * + num_days_between_for_dates('2014-01-05', '2014-01-15') + + ['2014-01-20'] * + num_days_between_for_dates('2014-01-16', '2014-01-20') + + ['NaT'] * + num_days_between_for_dates('2014-01-21', None) + ), + 1: zip_with_dates_for_dates( + ['NaT'] * + num_days_between_for_dates(None, '2014-01-04') + + ['2014-01-20'] * + num_days_between_for_dates('2014-01-05', '2014-01-09') + + ['2014-01-15'] * + num_days_between_for_dates('2014-01-10', '2014-01-15') + + ['2014-01-20'] * + num_days_between_for_dates('2014-01-16', '2014-01-20') + + ['NaT'] * + num_days_between_for_dates('2014-01-21', None) + ), + 2: zip_with_dates_for_dates( + ['NaT'] * + num_days_between_for_dates(None, '2014-01-04') + + ['2014-01-10'] * + num_days_between_for_dates('2014-01-05', '2014-01-10') + + ['NaT'] * + num_days_between_for_dates('2014-01-11', '2014-01-14') + + ['2014-01-20'] * + num_days_between_for_dates('2014-01-15', '2014-01-20') + + ['NaT'] * + num_days_between_for_dates('2014-01-21', None) + ), + 3: zip_with_dates_for_dates( + ['NaT'] * + num_days_between_for_dates(None, '2014-01-04') + + ['2014-01-10'] * + num_days_between_for_dates('2014-01-05', '2014-01-10') + + ['2014-01-15'] * + num_days_between_for_dates('2014-01-11', '2014-01-15') + + ['NaT'] * + num_days_between_for_dates('2014-01-16', None) + ), + 4: zip_with_dates_for_dates(['NaT'] * + len(dates)), + }, index=dates) + + def get_expected_previous_event_dates(self, dates): num_days_between_for_dates = partial(self.num_days_between, dates) zip_with_dates_for_dates = partial(self.zip_with_dates, dates) return pd.DataFrame({ @@ -339,7 +399,7 @@ class EventLoaderCommonTest(object): index=announcement_dates.index, ) - def _test_compute_buyback_auth(self, dates): + def _test_compute(self, dates): engine = self.setup_engine(dates) self.setup(dates) @@ -358,7 +418,3 @@ class EventLoaderCommonTest(object): assert_series_equal(result[col_name].xs(sid, level=1), self.cols[col_name][sid], check_names=False) - - - -