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

Code before:
            algo.future_symbol([1])

        with self.assertRaises(TypeError):
            algo.future_symbol({'foo': 'bar'})

    def test_future_chain_offset(self):
        algo = TradingAlgorithm(env=self.env)
        algo.datetime = pd.Timestamp('2006-12-01', tz='UTC')

        self.assertEqual(
            algo.future_chain('CL', offset=1).as_of_date,
            self.trading_calendar.next_session_label(
                algo.datetime
            )
        )

        self.assertEqual(
            algo.future_chain("CL", offset=5).as_of_date,
            pd.Timestamp("2006-12-08", tz='UTC')
        )

        self.assertEqual(
            algo.future_chain("CL", offset=-10).as_of_date,
            pd.Timestamp("2006-11-16", tz='UTC')
        )

        self.assertEqual(
            algo.future_chain(
                "CL",
                as_of_date=pd.Timestamp("2016-08-31", tz='UTC'),
                offset=10
            ).as_of_date,
            pd.Timestamp("2016-09-15", tz='UTC')
        )

    def test_future_chain(self):
        algo = TradingAlgorithm(env=self.env)
        algo.datetime = pd.Timestamp('2006-12-01', tz='UTC')

        # Check that the fields of the FutureChain object are set correctly
        cl = algo.future_chain('CL')
        self.assertEqual(cl.root_symbol, 'CL')
        self.assertEqual(cl.as_of_date, algo.datetime)

        # Check the fields are set correctly if an as_of_date is supplied
        as_of_date = pd.Timestamp('1952-08-11', tz='UTC')

        cl = algo.future_chain('CL', as_of_date=as_of_date)
        self.assertEqual(cl.root_symbol, 'CL')
        self.assertEqual(cl.as_of_date, as_of_date)


Review comment (suggestion):
It's a little confusing that `offset=1` is implemented in terms of the calendar, but `offset=5` is hard-coded.

Thoughts on either doing `self.trading_calendar.sessions_window(algo.datetime, 5)[-1]`, or iterating calls to `{next,previous}_session_label`.

e.g.

```
In [20]: def iterate(func, n, arg):
    while n > 0:
        arg = func(arg)
        n -= 1
    return arg
   ....:

In [21]: iterate(cal.previous_session_label, 1, dt)
Out[21]: Timestamp('2016-08-29 00:00:00+0000', tz='UTC', offset='C')

In [22]: iterate(cal.previous_session_label, 5, dt)
Out[22]: Timestamp('2016-08-23 00:00:00+0000', tz='UTC', offset='C')
```


Code after:
            algo.future_symbol([1])

        with self.assertRaises(TypeError):
            algo.future_symbol({'foo': 'bar'})

    def test_future_chain_offset(self):
        # November 2006         December 2006
        # Su Mo Tu We Th Fr Sa  Su Mo Tu We Th Fr Sa
        #           1  2  3  4                  1  2
        #  5  6  7  8  9 10 11   3  4  5  6  7  8  9
        # 12 13 14 15 16 17 18  10 11 12 13 14 15 16
        # 19 20 21 22 23 24 25  17 18 19 20 21 22 23
        # 26 27 28 29 30        24 25 26 27 28 29 30
        #                       31

        algo = TradingAlgorithm(env=self.env)
        algo.datetime = pd.Timestamp('2006-12-01', tz='UTC')

        self.assertEqual(
            algo.future_chain('CL', offset=1).as_of_date,
            pd.Timestamp("2006-12-04", tz='UTC')
        )

        self.assertEqual(
            algo.future_chain("CL", offset=5).as_of_date,
            pd.Timestamp("2006-12-08", tz='UTC')
        )

        self.assertEqual(
            algo.future_chain("CL", offset=-10).as_of_date,
            pd.Timestamp("2006-11-16", tz='UTC')
        )

        # September 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
        self.assertEqual(
            algo.future_chain(
                "CL",
                as_of_date=pd.Timestamp("2016-08-31", tz='UTC'),
                offset=10
            ).as_of_date,
            pd.Timestamp("2016-09-15", tz='UTC')
        )

    def test_future_chain(self):
        algo = TradingAlgorithm(env=self.env)