Repository: quantopian/zipline
File: tests/test_tradesimulation.py (Python)

Code before:
                                   zf.WithTradingSessions,
                                   zf.ZiplineTestCase):

    ASSET_FINDER_EQUITY_SIDS = (1,)
    BENCHMARK_SID = 1
    # These dates are chosen to cross a DST transition.
    #      March 2016
    # Su Mo Tu We Th Fr Sa
    #        1  2  3  4  5
    #  6  7  8  9 10 11 12
    # 13 14 15 16 17 18 19
    # 20 21 22 23 24 25 26
    # 27 28 29 30 31
    START_DATE = pd.Timestamp('2016-03-10', tz='UTC')
    END_DATE = pd.Timestamp('2016-03-15', tz='UTC')

    @parameter_space(
        num_sessions=[1, 2, 3],
        data_frequency=['daily', 'minute'],
        emission_rate=['daily', 'minute'],
        __fail_fast=True,
    )
    def test_before_trading_start_runs_at_8_45(self,
                                               num_sessions,
                                               data_frequency,
                                               emission_rate):
        bts_times = []

        def initialize(algo, data):
            pass

        def before_trading_start(algo, data):
            bts_times.append(algo.get_datetime())

        sim_params = SimulationParameters(
            start_session=self.nyse_sessions[1],
            end_session=self.nyse_sessions[num_sessions],
            data_frequency=data_frequency,
            emission_rate=emission_rate,
            trading_calendar=self.trading_calendar,
        )

        self.run_algorithm(
            before_trading_start=before_trading_start,
            sim_params=sim_params,
        )

        self.assertEqual(len(bts_times), num_sessions)
        expected_times = [
            pd.Timestamp('2016-03-11 8:45', tz='US/Eastern').tz_convert('UTC'),
            pd.Timestamp('2016-03-14 8:45', tz='US/Eastern').tz_convert('UTC'),

Review comment (suggestion):
I'm not sure what this test was originally supposed to test, but it looks like for a while it's just beeing testing the length of `sim_params.sessions` in a very confusing way...

Code after:
    # Su Mo Tu We Th Fr Sa
    #        1  2  3  4  5
    #  6  7  8  9 10 11 12
    # 13 14 15 16 17 18 19
    # 20 21 22 23 24 25 26
    # 27 28 29 30 31
    START_DATE = pd.Timestamp('2016-03-10', tz='UTC')
    END_DATE = pd.Timestamp('2016-03-15', tz='UTC')

    @parameter_space(
        num_sessions=[1, 2, 3],
        data_frequency=['daily', 'minute'],
        emission_rate=['daily', 'minute'],
        __fail_fast=True,
    )
    def test_before_trading_start_runs_at_8_45(self,
                                               num_sessions,
                                               data_frequency,
                                               emission_rate):
        bts_times = []

        def initialize(algo, data):
            pass

        def before_trading_start(algo, data):
            bts_times.append(algo.get_datetime())

        sim_params = SimulationParameters(
            # start at index 1 so we have an extra day to calculate benchmark
            # returns.
            start_session=self.nyse_sessions[1],
            end_session=self.nyse_sessions[num_sessions],
            data_frequency=data_frequency,
            emission_rate=emission_rate,
            trading_calendar=self.trading_calendar,
        )

        self.run_algorithm(
            before_trading_start=before_trading_start,
            sim_params=sim_params,
        )

        self.assertEqual(len(bts_times), num_sessions)
        expected_times = [
            pd.Timestamp('2016-03-11 8:45', tz='US/Eastern').tz_convert('UTC'),
            pd.Timestamp('2016-03-14 8:45', tz='US/Eastern').tz_convert('UTC'),
            pd.Timestamp('2016-03-15 8:45', tz='US/Eastern').tz_convert('UTC'),
        ]
        self.assertEqual(bts_times, expected_times[:num_sessions])