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

Code before:
        )

    def fake_pipeline(self):
        raise AssertionError

    def test_pipeline_compute_before_bts(self):

        # for storing and keeping track of calls to BTS and TestFactor.compute
        trace = []

        class TestFactor(CustomFactor):
            inputs = ()

            # window_length doesn't actually matter for this test case
            window_length = 1

            def compute(self, today, assets, out):
                trace.append("CustomFactor call")

        def initialize(context):
            pipeline = attach_pipeline(Pipeline(), 'my_pipeline')
            test_factor = TestFactor()
            pipeline.add(test_factor, 'test_factor')

        def before_trading_start(context, data):
            pipeline_output('my_pipeline')
            trace.append("BTS call")

        self.run_algorithm(
            initialize=initialize,
            before_trading_start=before_trading_start,
            sim_params=SimulationParameters(
                start_session=pd.Timestamp('2014-12-29', tz='utc'),
                end_session=pd.Timestamp('2014-12-31', tz='utc'),
                data_frequency='daily',
                emission_rate='daily',
                trading_calendar=self.trading_calendar,
            )
        )

        # All pipeline computation calls should occur before any BTS calls,
        # and the algorithm is being run for 3 days, so the first 3 calls
        # should be to the custom factor and the next 3 calls should be to BTS
        expected_result = ["CustomFactor call"] * 3 + ["BTS call"] * 3
        self.assertEqual(trace, expected_result)


Review comment (suggestion):
probably worth swapping the `pipeline_output` with the `trace` so that "BTS call" will appear first in the trace if there is an error.

Code after:
        )

    def fake_pipeline(self):
        raise AssertionError

    def test_pipeline_compute_before_bts(self):

        # for storing and keeping track of calls to BTS and TestFactor.compute
        trace = []

        class TestFactor(CustomFactor):
            inputs = ()

            # window_length doesn't actually matter for this test case
            window_length = 1

            def compute(self, today, assets, out):
                trace.append("CustomFactor call")

        def initialize(context):
            pipeline = attach_pipeline(Pipeline(), 'my_pipeline')
            test_factor = TestFactor()
            pipeline.add(test_factor, 'test_factor')

        def before_trading_start(context, data):
            trace.append("BTS call")
            pipeline_output('my_pipeline')

        self.run_algorithm(
            initialize=initialize,
            before_trading_start=before_trading_start,
            sim_params=SimulationParameters(
                start_session=pd.Timestamp('2014-12-29', tz='utc'),
                end_session=pd.Timestamp('2014-12-31', tz='utc'),
                data_frequency='daily',
                emission_rate='daily',
                trading_calendar=self.trading_calendar,
            )
        )

        # All pipeline computation calls should occur before any BTS calls,
        # and the algorithm is being run for 3 days, so the first 3 calls
        # should be to the custom factor and the next 3 calls should be to BTS
        expected_result = ["CustomFactor call"] * 3 + ["BTS call"] * 3
        self.assertEqual(trace, expected_result)