TEST: Simplify testing of restriction orderings.

This commit is contained in:
Scott Sanderson
2016-10-05 14:42:17 -04:00
parent b0aba20a6e
commit 8465ea18d0
+19 -42
View File
@@ -60,60 +60,37 @@ class RestrictionsTestCase(WithDataPortal, ZiplineTestCase):
pd.Timedelta('1 minute'),
pd.Timedelta('15 hours 5 minutes')
),
check_unordered=(False, True),
check_mixed=(False, True),
restriction_order=(
list(range(6)), # Keep restrictions in order.
[0, 2, 1, 3, 5, 4], # Re-order within asset.
[0, 3, 1, 4, 2, 5], # Scramble assets, maintain per-asset order.
[0, 5, 2, 3, 1, 4], # Scramble assets and per-asset order.
),
__fail_fast=True,
)
def test_historical_restrictions(self, date_offset, check_unordered,
check_mixed):
def test_historical_restrictions(self, date_offset, restriction_order):
"""
Test historical restrictions for both interday and intraday
restrictions, as well as restrictions defined in/not in order, for both
single- and multi-asset queries
"""
if check_unordered:
def maybe_scramble(rs):
# Swap the first two restrictions to check that we don't care
# that the restriction dates are ordered.
tmp = rs[0]
rs[0] = rs[1]
rs[1] = tmp
return rs
else:
maybe_scramble = lambda r: r
if check_mixed:
def maybe_mix_assets(rl):
# In the final restrictions list containing restrictions of
# both assets, swap the first and last restrictions so that
# the restrictions are not contiguous by asset
tmp = rl[0]
rl[0] = rl[-1]
rl[-1] = tmp
return rl
else:
def maybe_mix_assets(rl):
return rl
def rdate(s):
"""Convert a date string into a restriction for that date."""
# Add date_offset to check that we handle intraday changes.
return str_to_ts(s) + date_offset
all_restrictions = (
maybe_scramble([
Restriction(self.ASSET1, rdate('2011-01-04'), FROZEN),
Restriction(self.ASSET1, rdate('2011-01-05'), ALLOWED),
Restriction(self.ASSET1, rdate('2011-01-06'), FROZEN),
])
+
maybe_scramble([
Restriction(self.ASSET2, rdate('2011-01-05'), FROZEN),
Restriction(self.ASSET2, rdate('2011-01-06'), ALLOWED),
Restriction(self.ASSET2, rdate('2011-01-07'), FROZEN),
])
)
all_restrictions = maybe_mix_assets(all_restrictions)
base_restrictions = [
Restriction(self.ASSET1, rdate('2011-01-04'), FROZEN),
Restriction(self.ASSET1, rdate('2011-01-05'), ALLOWED),
Restriction(self.ASSET1, rdate('2011-01-06'), FROZEN),
Restriction(self.ASSET2, rdate('2011-01-05'), FROZEN),
Restriction(self.ASSET2, rdate('2011-01-06'), ALLOWED),
Restriction(self.ASSET2, rdate('2011-01-07'), FROZEN),
]
# Scramble the restrictions based on restriction_order to check that we
# don't depend on the order in which restrictions are provided to us.
all_restrictions = [base_restrictions[i] for i in restriction_order]
restrictions_by_asset = groupby(lambda r: r.asset, all_restrictions)
rl = HistoricalRestrictions(all_restrictions)