diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 9050086e..491df1b5 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1,20 +1,152 @@ -from datetime import timedelta +import pytz + +from datetime import timedelta, datetime from collections import defaultdict from unittest2 import TestCase -from zipline.utils.test_utils import setup_logger, teardown_logger +from zipline import ndict -import zipline.utils.factory as factory +from zipline.lines import SimulatedTrading + +from zipline.utils.test_utils import setup_logger, teardown_logger +from zipline.utils.date_utils import utcnow from zipline.gens.tradegens import SpecificEquityTrades -from zipline.gens.transform import StatefulTransform +from zipline.gens.transform import StatefulTransform, EventWindow from zipline.gens.vwap import VWAP from zipline.gens.mavg import MovingAverage from zipline.gens.returns import Returns -from zipline.lines import SimulatedTrading -from zipline.core.devsimulator import AddressAllocator -allocator = AddressAllocator(1000) +import zipline.utils.factory as factory + +def to_dt(msg): + return ndict({'dt': msg}) + +class NoopEventWindow(EventWindow): + """ + A no-op EventWindow subclass for testing the base EventWindow logic. + Keeps lists of all added and dropped events. + """ + def __init__(self, market_aware, days, delta): + EventWindow.__init__(self, market_aware, days, delta) + + self.added = [] + self.removed = [] + + def handle_add(self, event): + self.added.append(event) + + def handle_remove(self, event): + self.removed.append(event) + +class EventWindowTestCase(TestCase): + + def setUp(self): + setup_logger(self) + + # Constants calling before open, during the day, and after + # close on a valid trading day. + self.pre_open = datetime(2012, 8, 7, 13, tzinfo = pytz.utc) + self.mid_day = datetime(2012, 8, 7, 15, tzinfo = pytz.utc) + self.post_close = datetime(2012, 8, 7, 22, tzinfo = pytz.utc) + + # Constants calling before open, during the day, and after + # close on a saturday. + self.pre_open_saturday = datetime(2012, 8, 11, 13, tzinfo = pytz.utc) + self.mid_day_saturday = datetime(2012, 8, 11, 15, tzinfo = pytz.utc) + self.post_close_saturday = datetime(2012, 8, 11, 22, tzinfo = pytz.utc) + + # Constants calling before open, during the day, and after + # close on a holiday. + self.pre_open_holiday = datetime(2012, 12, 25, 13, tzinfo = pytz.utc) + self.mid_day_holiday = datetime(2012, 12, 25, tzinfo = pytz.utc) + self.post_close_holiday = datetime(2012, 12, 25, 22, tzinfo = pytz.utc) + + def test_event_window_with_timedelta(self): + + # Keep all events within a 5 minute window. + window = NoopEventWindow( + market_aware = False, + delta = timedelta(minutes = 5), + days = None + ) + now = utcnow() + + # 15 dates, increasing in 1 minute increments. + dates = [now + i * timedelta(minutes = 1) + for i in xrange(15)] + + # Turn the dates into the format required by EventWindow. + dt_messages = [to_dt(date) for date in dates] + + # Run all messages through the window and assert that we're adding + # and removing messages appropriately. We start the enumeration at 1 + # for convenience. + for num, message in enumerate(dt_messages, 1): + window.update(message) + + # Assert that we've added the correct number of events. + assert len(window.added) == num + + # Assert that we removed only events that fall outside (or + # on the boundary of) the delta. + for dropped in window.removed: + assert message.dt - dropped.dt >= timedelta(minutes = 5) + + def test_market_aware_window(self): + window = NoopEventWindow( + market_aware = True, + delta = None, + days = 1 + ) + dates = ([self.pre_open]*3) + dates += ([self.mid_day]*3) + dates += ([self.post_close]*3) + dates += [self.pre_open + timedelta(days = 1, seconds = 1)] + events = [to_dt(date) for date in dates] + + # Run the events. + for event in events: + window.update(event) + + # We should have removed the pre_open events on the first day. + # The rest should be intact. + + assert window.added == events + assert window.removed == events[0:3] + assert list(window.ticks) == events[3:] + + def test_market_aware_window_weekend(self): + window = NoopEventWindow( + market_aware = True, + delta = None, + days = 2 + ) + dates = [self.pre_open_saturday - timedelta(days = 1, seconds=1)] + dates += [self.mid_day_saturday - timedelta(days = 1, seconds=1)] + dates += [self.post_close_saturday - timedelta(days = 1, seconds=1)] + dates += [self.mid_day_saturday + timedelta(days = 1)] + + events = [to_dt(date) for date in dates] + + # Run the events. + for event in events: + window.update(event) + + # We shouldn't remove any events. + assert window.added == events + assert window.removed == [] + assert list(window.ticks) == events + + extra = to_dt(self.mid_day_saturday + timedelta(days = 2)) + window.update(extra) + + # We should remove only the first event. + assert window.removed == [events[0]] + assert list(window.ticks) == events[1:] + [extra] + + def tearDown(self): + setup_logger(self) class FinanceTransformsTestCase(TestCase): @@ -67,11 +199,10 @@ class FinanceTransformsTestCase(TestCase): # No returns for the first event because we don't have a # previous close. - expected = [None, 0.0, 0.1, 0.0] + expected = [0.0, 0.0, 0.1, 0.0] assert tnfm_vals == expected - - + # Two-day returns. An extra kink here is that the # factory will automatically skip a weekend for the # last event. Results shouldn't notice this blip. @@ -91,8 +222,8 @@ class FinanceTransformsTestCase(TestCase): tnfm_vals = [message.tnfm_value for message in transformed] expected = [ - None, - None, + 0.0, + 0.0, (13.0 - 10.0) / 10.0, (12.0 - 15.0) / 15.0, (13.0 - 13.0) / 13.0 diff --git a/zipline/gens/returns.py b/zipline/gens/returns.py index 055b9735..ccfc7b00 100644 --- a/zipline/gens/returns.py +++ b/zipline/gens/returns.py @@ -34,7 +34,7 @@ class ReturnsFromPriorClose(object): def __init__(self, days): self.closes = deque() self.last_event = None - self.returns = None + self.returns = 0.0 self.days = days def get_returns(self): diff --git a/zipline/gens/transform.py b/zipline/gens/transform.py index dbb7c75b..651c337d 100644 --- a/zipline/gens/transform.py +++ b/zipline/gens/transform.py @@ -125,7 +125,8 @@ class StatefulTransform(object): # returned by state.update(event). This is almost # identical to the behavior of FORWARDER, except we # compress the two calculated values (tnfm_id, and - # tnfm_value) into a single field. + # tnfm_value) into a single field. This mode is used by + # the sequential_transforms composite. elif self.append_value: out_message = message_copy out_message[self.namestring] = tnfm_value diff --git a/zipline/lines.py b/zipline/lines.py index 5fe87a0d..0fa990ca 100644 --- a/zipline/lines.py +++ b/zipline/lines.py @@ -307,7 +307,7 @@ class SimulatedTrading(object): simulation_style = SIMULATION_STYLE.FIXED_SLIPPAGE zmq_context = config.get('zmq_context', None) - simulation_id = config.get('simumlation_id', 'test_simulation') + simulation_id = config.get('simulation_id', 'test_simulation') results_socket_uri = config.get('results_socket_uri', None) #-------------------