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

Code before:
        self.check_downsampled_term(sma.quantiles(5))

    def test_errors_on_bad_downsample_frequency(self):

        f = NDaysAgoFactor(window_length=3)
        with self.assertRaises(ValueError) as e:
            f.downsample('bad')

        expected = (
            "{}() expected a value in "
            "('month_start', 'quarter_start', 'week_start', 'year_start') "
            "for argument 'frequency', but got 'bad' instead."
        ).format(_qualified_name(f.downsample))
        self.assertEqual(str(e.exception), expected)


class TestDownsampledRowwiseOperation(WithAssetFinder, ZiplineTestCase):
    T = partial(pd.Timestamp, tz='utc')
    START_DATE = T('2014-01-01')
    END_DATE = T('2014-02-01')
    HALF_WAY_POINT = T('2014-01-15')

    dates = pd.date_range(START_DATE, END_DATE)
    factor = TestingDataSet.float_col.latest

    @classmethod
    def init_class_fixtures(cls):
        super(TestDownsampledRowwiseOperation, cls).init_class_fixtures()

        def populate_initial_workspace(workspace,
                                       root_mask_term,
                                       execution_plan,
                                       dates,
                                       assets):
            """Fill the workspace value for ``cls.factor``.
            """
            # each value is just the sid
            base_row = np.array(sorted(cls.asset_finder.sids))

            # broadcast the row for each date
            workspace[cls.factor] = repeat_first_axis(base_row, len(dates))

            return workspace

        cls.pipeline_engine = SimplePipelineEngine(
            get_loader=lambda column: ExplodingObject(),
            calendar=cls.dates,
            asset_finder=cls.asset_finder,
            populate_initial_workspace=populate_initial_workspace,
        )


Review comment (suggestion):
Is there a particular reason that we want this test to use `populate_initial_workspace` to produce `factor`?  Another way we've produced this same factor in tests elsewhere has been to define a factor like:

```python
class Sid(CustomFactor):
    inputs = []
    window_length = 1
    def compute(self, today, assets, out):
        out[:] = assets
```

Code after:

    def test_errors_on_bad_downsample_frequency(self):

        f = NDaysAgoFactor(window_length=3)
        with self.assertRaises(ValueError) as e:
            f.downsample('bad')

        expected = (
            "{}() expected a value in "
            "('month_start', 'quarter_start', 'week_start', 'year_start') "
            "for argument 'frequency', but got 'bad' instead."
        ).format(_qualified_name(f.downsample))
        self.assertEqual(str(e.exception), expected)


class TestDownsampledRowwiseOperation(WithAssetFinder, ZiplineTestCase):

    T = partial(pd.Timestamp, tz='utc')
    START_DATE = T('2014-01-01')
    END_DATE = T('2014-02-01')
    HALF_WAY_POINT = T('2014-01-15')

    dates = pd.date_range(START_DATE, END_DATE)

    class SidFactor(CustomFactor):
        inputs = ()
        window_length = 1

        def compute(self, today, assets, out):
            out[:] = assets

    factor = SidFactor()

    @classmethod
    def init_class_fixtures(cls):
        super(TestDownsampledRowwiseOperation, cls).init_class_fixtures()
        cls.pipeline_engine = SimplePipelineEngine(
            get_loader=lambda c: ExplodingObject(),
            calendar=cls.dates,
            asset_finder=cls.asset_finder,
        )

    @classmethod
    def make_equity_info(cls):
        start = cls.START_DATE - pd.Timedelta(days=1)
        end = cls.END_DATE
        early_end = cls.HALF_WAY_POINT
        return pd.DataFrame(
            [['A',    'Ayy Inc.', start,       end, 'E', 'E'],
             ['B', 'early end',   start, early_end, 'E', 'E'],
             ['C',      'C Inc.', start,       end, 'E', 'E']],