TST: enhance test; add more common data.

STY: fixes for PEP8.
This commit is contained in:
Maya Tydykov
2016-02-22 14:09:39 -05:00
parent e257dc1da9
commit 5b37af6e04
3 changed files with 143 additions and 90 deletions
+16 -21
View File
@@ -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):
+34 -32
View File
@@ -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):
+93 -37
View File
@@ -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)