diff --git a/tests/test_tradesimulation.py b/tests/test_tradesimulation.py index d3a42fb6..d3bb8b33 100644 --- a/tests/test_tradesimulation.py +++ b/tests/test_tradesimulation.py @@ -21,8 +21,14 @@ from nose_parameterized import parameterized from six.moves import range from unittest import TestCase from zipline import TradingAlgorithm +from zipline.gens.sim_engine import BEFORE_TRADING_START_BAR + +from zipline.finance.performance import PerformanceTracker +from zipline.gens.tradesimulation import AlgorithmSimulator from zipline.sources.benchmark_source import BenchmarkSource from zipline.test_algorithms import NoopAlgorithm +from zipline.testing.fixtures import WithSimParams, ZiplineTestCase, \ + WithDataPortal from zipline.utils import factory from zipline.testing.core import FakeDataPortal from zipline.utils.calendars.trading_calendar import days_at_time @@ -93,3 +99,48 @@ class TestTradeSimulation(TestCase): ), "Expected %s but was %s." % (params.sessions, algo.before_trading_at)) + + +class BeforeTradingStartsOnlyClock(object): + def __init__(self, bts_minute): + self.bts_minute = bts_minute + + def __iter__(self): + yield self.bts_minute, BEFORE_TRADING_START_BAR + + +class TestBeforeTradingStartSimulationDt(WithSimParams, + WithDataPortal, + ZiplineTestCase): + + def test_bts_simulation_dt(self): + code = """ +def initialize(context): + pass +""" + algo = TradingAlgorithm(script=code, + sim_params=self.sim_params, + env=self.env) + + algo.perf_tracker = PerformanceTracker( + sim_params=self.sim_params, + trading_calendar=self.trading_calendar, + env=self.env, + ) + + dt = pd.Timestamp("2016-08-04 9:13:14", tz='US/Eastern') + algo_simulator = AlgorithmSimulator( + algo, + self.sim_params, + self.data_portal, + BeforeTradingStartsOnlyClock(dt), + algo._create_benchmark_source(), + None + ) + + # run through the algo's simulation + list(algo_simulator.transform()) + + # since the clock only ever emitted a single before_trading_start + # event, we can check that the simulation_dt was properly set + self.assertEqual(dt, algo_simulator.simulation_dt) diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index b8bbe2bb..31264828 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -229,7 +229,7 @@ class AlgorithmSimulator(object): yield self._get_daily_message(dt, algo, algo.perf_tracker) elif action == BEFORE_TRADING_START_BAR: - # call before trading start + self.simulation_dt = dt algo.on_dt_changed(dt) algo.before_trading_start(self.current_data) elif action == MINUTE_END: