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

Code before:
            missing_value='',
        )

        expected_data = array([
            ['aaa-foo', 'bbb-foo', 'ccc-foo'],
            ['ddd-foo', 'eee-foo', 'fff-foo'],
            ['ggg-foo', 'hhh-foo', 'iii-foo'],
            ['jjj-foo', 'kkk-foo', 'lll-foo'],
            ['mmm-foo', 'nnn-foo', 'ooo-foo'],
        ])
        expected_label_array = LabelArray(expected_data, missing_value='')

        expected_adj_array = AdjustedArray(
            data=expected_label_array,
            adjustments={4: [ObjectOverwrite(2, 3, 0, 0, 'ppp-foo')]},
            missing_value='',
        )

        adj_array.update_labels(lambda x: x + '-foo')

        # Check that the mapped AdjustedArray has the expected baseline
        # values and adjustment values.
        check_arrays(adj_array.data, expected_adj_array.data)
        self.assertEqual(adj_array.adjustments, expected_adj_array.adjustments)

    @parameterized.expand([
        # Append new adjustment to back of the list if index already contains
        # adjustments.
        (
            True,
            {
                1: [Float64Multiply(0, 3, 0, 0, 4.0),
                    Float64Overwrite(2, 3, 0, 0, 1.0)],
                3: [Float64Multiply(0, 3, 0, 0, 4.0)],
            },
        ),
        # Append new adjustment to front of the list if index already contains
        # adjustments.
        (
            False,
            {
                1: [Float64Overwrite(2, 3, 0, 0, 1.0),
                    Float64Multiply(0, 3, 0, 0, 4.0)],
                3: [Float64Multiply(0, 3, 0, 0, 4.0)],
            },
        )
    ])
    def test_append_adjustments(self, append_back, expected_adjustments):
        adjustments = {
            1: [Float64Multiply(0, 3, 0, 0, 4.0)],
        }

Review comment (suggestion):
This is a good start, but I think there are a few important cases we should probably cover:

1. Initial adjustments have a key not present in `adjustments_to_append`. It should still show up in the output.
2. Initial adjustments has more than one adjustment at a given index.
3. `adjustments_to_append` has more than one adjustment at a given index.


Code after:
            missing_value='',
        )

        expected_data = array([
            ['aaa-foo', 'bbb-foo', 'ccc-foo'],
            ['ddd-foo', 'eee-foo', 'fff-foo'],
            ['ggg-foo', 'hhh-foo', 'iii-foo'],
            ['jjj-foo', 'kkk-foo', 'lll-foo'],
            ['mmm-foo', 'nnn-foo', 'ooo-foo'],
        ])
        expected_label_array = LabelArray(expected_data, missing_value='')

        expected_adj_array = AdjustedArray(
            data=expected_label_array,
            adjustments={4: [ObjectOverwrite(2, 3, 0, 0, 'ppp-foo')]},
            missing_value='',
        )

        adj_array.update_labels(lambda x: x + '-foo')

        # Check that the mapped AdjustedArray has the expected baseline
        # values and adjustment values.
        check_arrays(adj_array.data, expected_adj_array.data)
        self.assertEqual(adj_array.adjustments, expected_adj_array.adjustments)

    A = Float64Multiply(0, 4, 1, 1, 0.5)
    B = Float64Overwrite(3, 3, 4, 4, 4.2)
    C = Float64Multiply(0, 2, 0, 0, 0.14)
    D = Float64Overwrite(0, 3, 0, 0, 4.0)
    E = Float64Overwrite(0, 0, 1, 1, 3.7)
    F = Float64Multiply(0, 4, 3, 3, 10.0)
    G = Float64Overwrite(5, 5, 4, 4, 1.7)
    H = Float64Multiply(0, 4, 2, 2, 0.99)
    S = Float64Multiply(0, 1, 4, 4, 5.06)

    @parameterized.expand([(
        # Initial adjustments
        {
            1: [A, B],
            2: [C],
            4: [D],
        },

        # Adjustments to add
        {
            1: [E],
            2: [F, G],
            3: [H, S],
        },

        # Expected adjustments with 'append'