Completed
Pull Request — master (#941)
by James
01:20
created

tests.TestFuture.test_lookup_future_symbol()   B

Complexity

Conditions 5

Size

Total Lines 34

Duplication

Lines 0
Ratio 0 %
Metric Value
cc 5
dl 0
loc 34
rs 8.0896
1
#
2
# Copyright 2015 Quantopian, Inc.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15
16
"""
17
Tests for the zipline.assets package
18
"""
19
from contextlib import contextmanager
20
from datetime import datetime, timedelta
21
import pickle
22
import sys
23
from unittest import TestCase
24
import uuid
25
import warnings
26
27
import pandas as pd
28
from pandas.tseries.tools import normalize_date
29
from pandas.util.testing import assert_frame_equal
30
31
from nose_parameterized import parameterized
32
from numpy import full
33
import sqlalchemy as sa
34
35
from zipline.assets import (
36
    Asset,
37
    Equity,
38
    Future,
39
    AssetFinder,
40
    AssetFinderCachedEquities,
41
)
42
from six import itervalues
43
from toolz import valmap
44
45
from zipline.assets.futures import (
46
    cme_code_to_month,
47
    FutureChain,
48
    month_to_cme_code
49
)
50
from zipline.assets.asset_writer import (
51
    check_version_info,
52
    write_version_info,
53
)
54
from zipline.assets.asset_db_schema import (
55
    ASSET_DB_VERSION,
56
    _version_table_schema,
57
)
58
from zipline.errors import (
59
    EquitiesNotFound,
60
    FutureContractsNotFound,
61
    MultipleSymbolsFound,
62
    RootSymbolNotFound,
63
    AssetDBVersionError,
64
    SidAssignmentError,
65
    SidsNotFound,
66
    SymbolNotFound,
67
)
68
from zipline.finance.trading import TradingEnvironment, noop_load
69
from zipline.utils.test_utils import (
70
    all_subindices,
71
    make_commodity_future_info,
72
    make_rotating_equity_info,
73
    make_simple_equity_info,
74
    tmp_assets_db,
75
    tmp_asset_finder,
76
)
77
78
79
@contextmanager
80
def build_lookup_generic_cases(asset_finder_type):
81
    """
82
    Generate test cases for the type of asset finder specific by
83
    asset_finder_type for test_lookup_generic.
84
    """
85
86
    unique_start = pd.Timestamp('2013-01-01', tz='UTC')
87
    unique_end = pd.Timestamp('2014-01-01', tz='UTC')
88
89
    dupe_0_start = pd.Timestamp('2013-01-01', tz='UTC')
90
    dupe_0_end = dupe_0_start + timedelta(days=1)
91
92
    dupe_1_start = pd.Timestamp('2013-01-03', tz='UTC')
93
    dupe_1_end = dupe_1_start + timedelta(days=1)
94
95
    frame = pd.DataFrame.from_records(
96
        [
97
            {
98
                'sid': 0,
99
                'symbol': 'duplicated',
100
                'start_date': dupe_0_start.value,
101
                'end_date': dupe_0_end.value,
102
                'exchange': '',
103
            },
104
            {
105
                'sid': 1,
106
                'symbol': 'duplicated',
107
                'start_date': dupe_1_start.value,
108
                'end_date': dupe_1_end.value,
109
                'exchange': '',
110
            },
111
            {
112
                'sid': 2,
113
                'symbol': 'unique',
114
                'start_date': unique_start.value,
115
                'end_date': unique_end.value,
116
                'exchange': '',
117
            },
118
        ],
119
        index='sid')
120
    with tmp_assets_db(equities=frame) as assets_db:
121
        finder = asset_finder_type(assets_db)
122
        dupe_0, dupe_1, unique = assets = [
123
            finder.retrieve_asset(i)
124
            for i in range(3)
125
        ]
126
127
        dupe_0_start = dupe_0.start_date
128
        dupe_1_start = dupe_1.start_date
129
        yield (
130
            ##
131
            # Scalars
132
133
            # Asset object
134
            (finder, assets[0], None, assets[0]),
135
            (finder, assets[1], None, assets[1]),
136
            (finder, assets[2], None, assets[2]),
137
            # int
138
            (finder, 0, None, assets[0]),
139
            (finder, 1, None, assets[1]),
140
            (finder, 2, None, assets[2]),
141
            # Duplicated symbol with resolution date
142
            (finder, 'DUPLICATED', dupe_0_start, dupe_0),
143
            (finder, 'DUPLICATED', dupe_1_start, dupe_1),
144
            # Unique symbol, with or without resolution date.
145
            (finder, 'UNIQUE', unique_start, unique),
146
            (finder, 'UNIQUE', None, unique),
147
148
            ##
149
            # Iterables
150
151
            # Iterables of Asset objects.
152
            (finder, assets, None, assets),
153
            (finder, iter(assets), None, assets),
154
            # Iterables of ints
155
            (finder, (0, 1), None, assets[:-1]),
156
            (finder, iter((0, 1)), None, assets[:-1]),
157
            # Iterables of symbols.
158
            (finder, ('DUPLICATED', 'UNIQUE'), dupe_0_start, [dupe_0, unique]),
159
            (finder, ('DUPLICATED', 'UNIQUE'), dupe_1_start, [dupe_1, unique]),
160
            # Mixed types
161
            (finder,
162
             ('DUPLICATED', 2, 'UNIQUE', 1, dupe_1),
163
             dupe_0_start,
164
             [dupe_0, assets[2], unique, assets[1], dupe_1]),
165
        )
166
167
168
class AssetTestCase(TestCase):
169
170
    def test_asset_object(self):
171
        self.assertEquals({5061: 'foo'}[Asset(5061)], 'foo')
172
        self.assertEquals(Asset(5061), 5061)
173
        self.assertEquals(5061, Asset(5061))
174
175
        self.assertEquals(Asset(5061), Asset(5061))
176
        self.assertEquals(int(Asset(5061)), 5061)
177
178
        self.assertEquals(str(Asset(5061)), 'Asset(5061)')
179
180
    def test_asset_is_pickleable(self):
181
182
        # Very wow
183
        s = Asset(
184
            1337,
185
            symbol="DOGE",
186
            asset_name="DOGECOIN",
187
            start_date=pd.Timestamp('2013-12-08 9:31AM', tz='UTC'),
188
            end_date=pd.Timestamp('2014-06-25 11:21AM', tz='UTC'),
189
            first_traded=pd.Timestamp('2013-12-08 9:31AM', tz='UTC'),
190
            exchange='THE MOON',
191
        )
192
        s_unpickled = pickle.loads(pickle.dumps(s))
193
194
        attrs_to_check = ['end_date',
195
                          'exchange',
196
                          'first_traded',
197
                          'end_date',
198
                          'asset_name',
199
                          'start_date',
200
                          'sid',
201
                          'start_date',
202
                          'symbol']
203
204
        for attr in attrs_to_check:
205
            self.assertEqual(getattr(s, attr), getattr(s_unpickled, attr))
206
207
    def test_asset_comparisons(self):
208
209
        s_23 = Asset(23)
210
        s_24 = Asset(24)
211
212
        self.assertEqual(s_23, s_23)
213
        self.assertEqual(s_23, 23)
214
        self.assertEqual(23, s_23)
215
216
        self.assertNotEqual(s_23, s_24)
217
        self.assertNotEqual(s_23, 24)
218
        self.assertNotEqual(s_23, "23")
219
        self.assertNotEqual(s_23, 23.5)
220
        self.assertNotEqual(s_23, [])
221
        self.assertNotEqual(s_23, None)
222
223
        self.assertLess(s_23, s_24)
224
        self.assertLess(s_23, 24)
225
        self.assertGreater(24, s_23)
226
        self.assertGreater(s_24, s_23)
227
228
    def test_lt(self):
229
        self.assertTrue(Asset(3) < Asset(4))
230
        self.assertFalse(Asset(4) < Asset(4))
231
        self.assertFalse(Asset(5) < Asset(4))
232
233
    def test_le(self):
234
        self.assertTrue(Asset(3) <= Asset(4))
235
        self.assertTrue(Asset(4) <= Asset(4))
236
        self.assertFalse(Asset(5) <= Asset(4))
237
238
    def test_eq(self):
239
        self.assertFalse(Asset(3) == Asset(4))
240
        self.assertTrue(Asset(4) == Asset(4))
241
        self.assertFalse(Asset(5) == Asset(4))
242
243
    def test_ge(self):
244
        self.assertFalse(Asset(3) >= Asset(4))
245
        self.assertTrue(Asset(4) >= Asset(4))
246
        self.assertTrue(Asset(5) >= Asset(4))
247
248
    def test_gt(self):
249
        self.assertFalse(Asset(3) > Asset(4))
250
        self.assertFalse(Asset(4) > Asset(4))
251
        self.assertTrue(Asset(5) > Asset(4))
252
253
    def test_type_mismatch(self):
254
        if sys.version_info.major < 3:
255
            self.assertIsNotNone(Asset(3) < 'a')
256
            self.assertIsNotNone('a' < Asset(3))
257
        else:
258
            with self.assertRaises(TypeError):
259
                Asset(3) < 'a'
260
            with self.assertRaises(TypeError):
261
                'a' < Asset(3)
262
263
264
class TestFuture(TestCase):
265
266
    @classmethod
267
    def setUpClass(cls):
268
        cls.future = Future(
269
            2468,
270
            symbol='OMH15',
271
            root_symbol='OM',
272
            notice_date=pd.Timestamp('2014-01-20', tz='UTC'),
273
            expiration_date=pd.Timestamp('2014-02-20', tz='UTC'),
274
            auto_close_date=pd.Timestamp('2014-01-18', tz='UTC'),
275
            tick_size=.01,
276
            multiplier=500
277
        )
278
        cls.future2 = Future(
279
            0,
280
            symbol='CLG06',
281
            root_symbol='CL',
282
            start_date=pd.Timestamp('2005-12-01', tz='UTC'),
283
            notice_date=pd.Timestamp('2005-12-20', tz='UTC'),
284
            expiration_date=pd.Timestamp('2006-01-20', tz='UTC')
285
        )
286
        env = TradingEnvironment(load=noop_load)
287
        env.write_data(futures_identifiers=[TestFuture.future,
288
                                            TestFuture.future2])
289
        cls.asset_finder = env.asset_finder
290
291
    def test_str(self):
292
        strd = self.future.__str__()
293
        self.assertEqual("Future(2468 [OMH15])", strd)
294
295
    def test_repr(self):
296
        reprd = self.future.__repr__()
297
        self.assertTrue("Future" in reprd)
298
        self.assertTrue("2468" in reprd)
299
        self.assertTrue("OMH15" in reprd)
300
        self.assertTrue("root_symbol='OM'" in reprd)
301
        self.assertTrue(("notice_date=Timestamp('2014-01-20 00:00:00+0000', "
302
                        "tz='UTC')") in reprd)
303
        self.assertTrue("expiration_date=Timestamp('2014-02-20 00:00:00+0000'"
304
                        in reprd)
305
        self.assertTrue("auto_close_date=Timestamp('2014-01-18 00:00:00+0000'"
306
                        in reprd)
307
        self.assertTrue("tick_size=0.01" in reprd)
308
        self.assertTrue("multiplier=500" in reprd)
309
310
    def test_reduce(self):
311
        reduced = self.future.__reduce__()
312
        self.assertEqual(Future, reduced[0])
313
314
    def test_to_and_from_dict(self):
315
        dictd = self.future.to_dict()
316
        self.assertTrue('root_symbol' in dictd)
317
        self.assertTrue('notice_date' in dictd)
318
        self.assertTrue('expiration_date' in dictd)
319
        self.assertTrue('auto_close_date' in dictd)
320
        self.assertTrue('tick_size' in dictd)
321
        self.assertTrue('multiplier' in dictd)
322
323
        from_dict = Future.from_dict(dictd)
324
        self.assertTrue(isinstance(from_dict, Future))
325
        self.assertEqual(self.future, from_dict)
326
327
    def test_root_symbol(self):
328
        self.assertEqual('OM', self.future.root_symbol)
329
330
    def test_lookup_future_symbol(self):
331
        """
332
        Test the lookup_future_symbol method.
333
        """
334
        om = TestFuture.asset_finder.lookup_future_symbol('OMH15')
335
        self.assertEqual(om.sid, 2468)
336
        self.assertEqual(om.symbol, 'OMH15')
337
        self.assertEqual(om.root_symbol, 'OM')
338
        self.assertEqual(om.notice_date, pd.Timestamp('2014-01-20', tz='UTC'))
339
        self.assertEqual(om.expiration_date,
340
                         pd.Timestamp('2014-02-20', tz='UTC'))
341
        self.assertEqual(om.auto_close_date,
342
                         pd.Timestamp('2014-01-18', tz='UTC'))
343
344
        cl = TestFuture.asset_finder.lookup_future_symbol('CLG06')
345
        self.assertEqual(cl.sid, 0)
346
        self.assertEqual(cl.symbol, 'CLG06')
347
        self.assertEqual(cl.root_symbol, 'CL')
348
        self.assertEqual(cl.start_date, pd.Timestamp('2005-12-01', tz='UTC'))
349
        self.assertEqual(cl.notice_date, pd.Timestamp('2005-12-20', tz='UTC'))
350
        self.assertEqual(cl.expiration_date,
351
                         pd.Timestamp('2006-01-20', tz='UTC'))
352
353
        with self.assertRaises(SymbolNotFound):
354
            TestFuture.asset_finder.lookup_future_symbol('')
355
356
        with self.assertRaises(SymbolNotFound):
357
            TestFuture.asset_finder.lookup_future_symbol('#&?!')
358
359
        with self.assertRaises(SymbolNotFound):
360
            TestFuture.asset_finder.lookup_future_symbol('FOOBAR')
361
362
        with self.assertRaises(SymbolNotFound):
363
            TestFuture.asset_finder.lookup_future_symbol('XXX99')
364
365
366
class AssetFinderTestCase(TestCase):
367
368
    def setUp(self):
369
        self.env = TradingEnvironment(load=noop_load)
370
        self.asset_finder_type = AssetFinder
371
372
    def test_lookup_symbol_delimited(self):
373
        as_of = pd.Timestamp('2013-01-01', tz='UTC')
374
        frame = pd.DataFrame.from_records(
375
            [
376
                {
377
                    'sid': i,
378
                    'symbol':  'TEST.%d' % i,
379
                    'company_name': "company%d" % i,
380
                    'start_date': as_of.value,
381
                    'end_date': as_of.value,
382
                    'exchange': uuid.uuid4().hex
383
                }
384
                for i in range(3)
385
            ]
386
        )
387
        self.env.write_data(equities_df=frame)
388
        finder = self.asset_finder_type(self.env.engine)
389
        asset_0, asset_1, asset_2 = (
390
            finder.retrieve_asset(i) for i in range(3)
391
        )
392
393
        # we do it twice to catch caching bugs
394
        for i in range(2):
395
            with self.assertRaises(SymbolNotFound):
396
                finder.lookup_symbol('TEST', as_of)
397
            with self.assertRaises(SymbolNotFound):
398
                finder.lookup_symbol('TEST1', as_of)
399
            # '@' is not a supported delimiter
400
            with self.assertRaises(SymbolNotFound):
401
                finder.lookup_symbol('TEST@1', as_of)
402
403
            # Adding an unnecessary fuzzy shouldn't matter.
404
            for fuzzy_char in ['-', '/', '_', '.']:
405
                self.assertEqual(
406
                    asset_1,
407
                    finder.lookup_symbol('TEST%s1' % fuzzy_char, as_of)
408
                )
409
410
    def test_lookup_symbol_fuzzy(self):
411
        metadata = {
412
            0: {'symbol': 'PRTY_HRD'},
413
            1: {'symbol': 'BRKA'},
414
            2: {'symbol': 'BRK_A'},
415
        }
416
        self.env.write_data(equities_data=metadata)
417
        finder = self.env.asset_finder
418
        dt = pd.Timestamp('2013-01-01', tz='UTC')
419
420
        # Try combos of looking up PRTYHRD with and without a time or fuzzy
421
        # Both non-fuzzys get no result
422
        with self.assertRaises(SymbolNotFound):
423
            finder.lookup_symbol('PRTYHRD', None)
424
        with self.assertRaises(SymbolNotFound):
425
            finder.lookup_symbol('PRTYHRD', dt)
426
        # Both fuzzys work
427
        self.assertEqual(0, finder.lookup_symbol('PRTYHRD', None, fuzzy=True))
428
        self.assertEqual(0, finder.lookup_symbol('PRTYHRD', dt, fuzzy=True))
429
430
        # Try combos of looking up PRTY_HRD, all returning sid 0
431
        self.assertEqual(0, finder.lookup_symbol('PRTY_HRD', None))
432
        self.assertEqual(0, finder.lookup_symbol('PRTY_HRD', dt))
433
        self.assertEqual(0, finder.lookup_symbol('PRTY_HRD', None, fuzzy=True))
434
        self.assertEqual(0, finder.lookup_symbol('PRTY_HRD', dt, fuzzy=True))
435
436
        # Try combos of looking up BRKA, all returning sid 1
437
        self.assertEqual(1, finder.lookup_symbol('BRKA', None))
438
        self.assertEqual(1, finder.lookup_symbol('BRKA', dt))
439
        self.assertEqual(1, finder.lookup_symbol('BRKA', None, fuzzy=True))
440
        self.assertEqual(1, finder.lookup_symbol('BRKA', dt, fuzzy=True))
441
442
        # Try combos of looking up BRK_A, all returning sid 2
443
        self.assertEqual(2, finder.lookup_symbol('BRK_A', None))
444
        self.assertEqual(2, finder.lookup_symbol('BRK_A', dt))
445
        self.assertEqual(2, finder.lookup_symbol('BRK_A', None, fuzzy=True))
446
        self.assertEqual(2, finder.lookup_symbol('BRK_A', dt, fuzzy=True))
447
448
    def test_lookup_symbol(self):
449
450
        # Incrementing by two so that start and end dates for each
451
        # generated Asset don't overlap (each Asset's end_date is the
452
        # day after its start date.)
453
        dates = pd.date_range('2013-01-01', freq='2D', periods=5, tz='UTC')
454
        df = pd.DataFrame.from_records(
455
            [
456
                {
457
                    'sid': i,
458
                    'symbol':  'existing',
459
                    'start_date': date.value,
460
                    'end_date': (date + timedelta(days=1)).value,
461
                    'exchange': 'NYSE',
462
                }
463
                for i, date in enumerate(dates)
464
            ]
465
        )
466
        self.env.write_data(equities_df=df)
467
        finder = self.asset_finder_type(self.env.engine)
468
        for _ in range(2):  # Run checks twice to test for caching bugs.
469
            with self.assertRaises(SymbolNotFound):
470
                finder.lookup_symbol('NON_EXISTING', dates[0])
471
472
            with self.assertRaises(MultipleSymbolsFound):
473
                finder.lookup_symbol('EXISTING', None)
474
475
            for i, date in enumerate(dates):
476
                # Verify that we correctly resolve multiple symbols using
477
                # the supplied date
478
                result = finder.lookup_symbol('EXISTING', date)
479
                self.assertEqual(result.symbol, 'EXISTING')
480
                self.assertEqual(result.sid, i)
481
482
    def test_lookup_symbol_from_multiple_valid(self):
483
        # This test asserts that we resolve conflicts in accordance with the
484
        # following rules when we have multiple assets holding the same symbol
485
        # at the same time:
486
487
        # If multiple SIDs exist for symbol S at time T, return the candidate
488
        # SID whose start_date is highest. (200 cases)
489
490
        # If multiple SIDs exist for symbol S at time T, the best candidate
491
        # SIDs share the highest start_date, return the SID with the highest
492
        # end_date. (34 cases)
493
494
        # It is the opinion of the author (ssanderson) that we should consider
495
        # this malformed input and fail here.  But this is the current indended
496
        # behavior of the code, and I accidentally broke it while refactoring.
497
        # These will serve as regression tests until the time comes that we
498
        # decide to enforce this as an error.
499
500
        # See https://github.com/quantopian/zipline/issues/837 for more
501
        # details.
502
503
        df = pd.DataFrame.from_records(
504
            [
505
                {
506
                    'sid': 1,
507
                    'symbol': 'multiple',
508
                    'start_date': pd.Timestamp('2010-01-01'),
509
                    'end_date': pd.Timestamp('2012-01-01'),
510
                    'exchange': 'NYSE'
511
                },
512
                # Same as asset 1, but with a later end date.
513
                {
514
                    'sid': 2,
515
                    'symbol': 'multiple',
516
                    'start_date': pd.Timestamp('2010-01-01'),
517
                    'end_date': pd.Timestamp('2013-01-01'),
518
                    'exchange': 'NYSE'
519
                },
520
                # Same as asset 1, but with a later start_date
521
                {
522
                    'sid': 3,
523
                    'symbol': 'multiple',
524
                    'start_date': pd.Timestamp('2011-01-01'),
525
                    'end_date': pd.Timestamp('2012-01-01'),
526
                    'exchange': 'NYSE'
527
                },
528
            ]
529
        )
530
531
        def check(expected_sid, date):
532
            result = finder.lookup_symbol(
533
                'MULTIPLE', date,
534
            )
535
            self.assertEqual(result.symbol, 'MULTIPLE')
536
            self.assertEqual(result.sid, expected_sid)
537
538
        with tmp_asset_finder(finder_cls=self.asset_finder_type,
539
                              equities=df) as finder:
540
            self.assertIsInstance(finder, self.asset_finder_type)
541
542
            # Sids 1 and 2 are eligible here.  We should get asset 2 because it
543
            # has the later end_date.
544
            check(2, pd.Timestamp('2010-12-31'))
545
546
            # Sids 1, 2, and 3 are eligible here.  We should get sid 3 because
547
            # it has a later start_date
548
            check(3, pd.Timestamp('2011-01-01'))
549
550
    def test_lookup_generic(self):
551
        """
552
        Ensure that lookup_generic works with various permutations of inputs.
553
        """
554
        with build_lookup_generic_cases(self.asset_finder_type) as cases:
555
            for finder, symbols, reference_date, expected in cases:
556
                results, missing = finder.lookup_generic(symbols,
557
                                                         reference_date)
558
                self.assertEqual(results, expected)
559
                self.assertEqual(missing, [])
560
561
    def test_lookup_generic_handle_missing(self):
562
        data = pd.DataFrame.from_records(
563
            [
564
                {
565
                    'sid': 0,
566
                    'symbol': 'real',
567
                    'start_date': pd.Timestamp('2013-1-1', tz='UTC'),
568
                    'end_date': pd.Timestamp('2014-1-1', tz='UTC'),
569
                    'exchange': '',
570
                },
571
                {
572
                    'sid': 1,
573
                    'symbol': 'also_real',
574
                    'start_date': pd.Timestamp('2013-1-1', tz='UTC'),
575
                    'end_date': pd.Timestamp('2014-1-1', tz='UTC'),
576
                    'exchange': '',
577
                },
578
                # Sid whose end date is before our query date.  We should
579
                # still correctly find it.
580
                {
581
                    'sid': 2,
582
                    'symbol': 'real_but_old',
583
                    'start_date': pd.Timestamp('2002-1-1', tz='UTC'),
584
                    'end_date': pd.Timestamp('2003-1-1', tz='UTC'),
585
                    'exchange': '',
586
                },
587
                # Sid whose start_date is **after** our query date.  We should
588
                # **not** find it.
589
                {
590
                    'sid': 3,
591
                    'symbol': 'real_but_in_the_future',
592
                    'start_date': pd.Timestamp('2014-1-1', tz='UTC'),
593
                    'end_date': pd.Timestamp('2020-1-1', tz='UTC'),
594
                    'exchange': 'THE FUTURE',
595
                },
596
            ]
597
        )
598
        self.env.write_data(equities_df=data)
599
        finder = self.asset_finder_type(self.env.engine)
600
        results, missing = finder.lookup_generic(
601
            ['REAL', 1, 'FAKE', 'REAL_BUT_OLD', 'REAL_BUT_IN_THE_FUTURE'],
602
            pd.Timestamp('2013-02-01', tz='UTC'),
603
        )
604
605
        self.assertEqual(len(results), 3)
606
        self.assertEqual(results[0].symbol, 'REAL')
607
        self.assertEqual(results[0].sid, 0)
608
        self.assertEqual(results[1].symbol, 'ALSO_REAL')
609
        self.assertEqual(results[1].sid, 1)
610
        self.assertEqual(results[2].symbol, 'REAL_BUT_OLD')
611
        self.assertEqual(results[2].sid, 2)
612
613
        self.assertEqual(len(missing), 2)
614
        self.assertEqual(missing[0], 'FAKE')
615
        self.assertEqual(missing[1], 'REAL_BUT_IN_THE_FUTURE')
616
617
    def test_insert_metadata(self):
618
        data = {0: {'start_date': '2014-01-01',
619
                    'end_date': '2015-01-01',
620
                    'symbol': "PLAY",
621
                    'foo_data': "FOO"}}
622
        self.env.write_data(equities_data=data)
623
        finder = self.asset_finder_type(self.env.engine)
624
        # Test proper insertion
625
        equity = finder.retrieve_asset(0)
626
        self.assertIsInstance(equity, Equity)
627
        self.assertEqual('PLAY', equity.symbol)
628
        self.assertEqual(pd.Timestamp('2015-01-01', tz='UTC'),
629
                         equity.end_date)
630
631
        # Test invalid field
632
        with self.assertRaises(AttributeError):
633
            equity.foo_data
634
635
    def test_consume_metadata(self):
636
637
        # Test dict consumption
638
        dict_to_consume = {0: {'symbol': 'PLAY'},
639
                           1: {'symbol': 'MSFT'}}
640
        self.env.write_data(equities_data=dict_to_consume)
641
        finder = self.asset_finder_type(self.env.engine)
642
643
        equity = finder.retrieve_asset(0)
644
        self.assertIsInstance(equity, Equity)
645
        self.assertEqual('PLAY', equity.symbol)
646
647
        # Test dataframe consumption
648
        df = pd.DataFrame(columns=['asset_name', 'exchange'], index=[0, 1])
649
        df['asset_name'][0] = "Dave'N'Busters"
650
        df['exchange'][0] = "NASDAQ"
651
        df['asset_name'][1] = "Microsoft"
652
        df['exchange'][1] = "NYSE"
653
        self.env = TradingEnvironment(load=noop_load)
654
        self.env.write_data(equities_df=df)
655
        finder = self.asset_finder_type(self.env.engine)
656
        self.assertEqual('NASDAQ', finder.retrieve_asset(0).exchange)
657
        self.assertEqual('Microsoft', finder.retrieve_asset(1).asset_name)
658
659
    def test_consume_asset_as_identifier(self):
660
        # Build some end dates
661
        eq_end = pd.Timestamp('2012-01-01', tz='UTC')
662
        fut_end = pd.Timestamp('2008-01-01', tz='UTC')
663
664
        # Build some simple Assets
665
        equity_asset = Equity(1, symbol="TESTEQ", end_date=eq_end)
666
        future_asset = Future(200, symbol="TESTFUT", end_date=fut_end)
667
668
        # Consume the Assets
669
        self.env.write_data(equities_identifiers=[equity_asset],
670
                            futures_identifiers=[future_asset])
671
        finder = self.asset_finder_type(self.env.engine)
672
673
        # Test equality with newly built Assets
674
        self.assertEqual(equity_asset, finder.retrieve_asset(1))
675
        self.assertEqual(future_asset, finder.retrieve_asset(200))
676
        self.assertEqual(eq_end, finder.retrieve_asset(1).end_date)
677
        self.assertEqual(fut_end, finder.retrieve_asset(200).end_date)
678
679
    def test_sid_assignment(self):
680
681
        # This metadata does not contain SIDs
682
        metadata = ['PLAY', 'MSFT']
683
684
        today = normalize_date(pd.Timestamp('2015-07-09', tz='UTC'))
685
686
        # Write data with sid assignment
687
        self.env.write_data(equities_identifiers=metadata,
688
                            allow_sid_assignment=True)
689
690
        # Verify that Assets were built and different sids were assigned
691
        finder = self.asset_finder_type(self.env.engine)
692
        play = finder.lookup_symbol('PLAY', today)
693
        msft = finder.lookup_symbol('MSFT', today)
694
        self.assertEqual('PLAY', play.symbol)
695
        self.assertIsNotNone(play.sid)
696
        self.assertNotEqual(play.sid, msft.sid)
697
698
    def test_sid_assignment_failure(self):
699
700
        # This metadata does not contain SIDs
701
        metadata = ['PLAY', 'MSFT']
702
703
        # Write data without sid assignment, asserting failure
704
        with self.assertRaises(SidAssignmentError):
705
            self.env.write_data(equities_identifiers=metadata,
706
                                allow_sid_assignment=False)
707
708
    def test_security_dates_warning(self):
709
710
        # Build an asset with an end_date
711
        eq_end = pd.Timestamp('2012-01-01', tz='UTC')
712
        equity_asset = Equity(1, symbol="TESTEQ", end_date=eq_end)
713
714
        # Catch all warnings
715
        with warnings.catch_warnings(record=True) as w:
716
            # Cause all warnings to always be triggered
717
            warnings.simplefilter("always")
718
            equity_asset.security_start_date
719
            equity_asset.security_end_date
720
            equity_asset.security_name
721
            # Verify the warning
722
            self.assertEqual(3, len(w))
723
            for warning in w:
724
                self.assertTrue(issubclass(warning.category,
725
                                           DeprecationWarning))
726
727
    def test_lookup_future_chain(self):
728
        metadata = {
729
            # Notice day is today, so should be valid.
730
            0: {
731
                'symbol': 'ADN15',
732
                'root_symbol': 'AD',
733
                'notice_date': pd.Timestamp('2015-05-14', tz='UTC'),
734
                'expiration_date': pd.Timestamp('2015-06-14', tz='UTC'),
735
                'start_date': pd.Timestamp('2015-01-01', tz='UTC')
736
            },
737
            1: {
738
                'symbol': 'ADV15',
739
                'root_symbol': 'AD',
740
                'notice_date': pd.Timestamp('2015-08-14', tz='UTC'),
741
                'expiration_date': pd.Timestamp('2015-09-14', tz='UTC'),
742
                'start_date': pd.Timestamp('2015-01-01', tz='UTC')
743
            },
744
            # Starts trading today, so should be valid.
745
            2: {
746
                'symbol': 'ADF16',
747
                'root_symbol': 'AD',
748
                'notice_date': pd.Timestamp('2015-11-16', tz='UTC'),
749
                'expiration_date': pd.Timestamp('2015-12-16', tz='UTC'),
750
                'start_date': pd.Timestamp('2015-05-14', tz='UTC')
751
            },
752
            # Starts trading in August, so not valid.
753
            3: {
754
                'symbol': 'ADX16',
755
                'root_symbol': 'AD',
756
                'notice_date': pd.Timestamp('2015-11-16', tz='UTC'),
757
                'expiration_date': pd.Timestamp('2015-12-16', tz='UTC'),
758
                'start_date': pd.Timestamp('2015-08-01', tz='UTC')
759
            },
760
            # Notice date comes after expiration
761
            4: {
762
                'symbol': 'ADZ16',
763
                'root_symbol': 'AD',
764
                'notice_date': pd.Timestamp('2016-11-25', tz='UTC'),
765
                'expiration_date': pd.Timestamp('2016-11-16', tz='UTC'),
766
                'start_date': pd.Timestamp('2015-08-01', tz='UTC')
767
            },
768
            # This contract has no start date and also this contract should be
769
            # last in all chains
770
            5: {
771
                'symbol': 'ADZ20',
772
                'root_symbol': 'AD',
773
                'notice_date': pd.Timestamp('2020-11-25', tz='UTC'),
774
                'expiration_date': pd.Timestamp('2020-11-16', tz='UTC')
775
            },
776
        }
777
        self.env.write_data(futures_data=metadata)
778
        finder = self.asset_finder_type(self.env.engine)
779
        dt = pd.Timestamp('2015-05-14', tz='UTC')
780
        dt_2 = pd.Timestamp('2015-10-14', tz='UTC')
781
        dt_3 = pd.Timestamp('2016-11-17', tz='UTC')
782
783
        # Check that we get the expected number of contracts, in the
784
        # right order
785
        ad_contracts = finder.lookup_future_chain('AD', dt)
786
        self.assertEqual(len(ad_contracts), 6)
787
        self.assertEqual(ad_contracts[0].sid, 0)
788
        self.assertEqual(ad_contracts[1].sid, 1)
789
        self.assertEqual(ad_contracts[5].sid, 5)
790
791
        # Check that, when some contracts have expired, the chain has advanced
792
        # properly to the next contracts
793
        ad_contracts = finder.lookup_future_chain('AD', dt_2)
794
        self.assertEqual(len(ad_contracts), 4)
795
        self.assertEqual(ad_contracts[0].sid, 2)
796
        self.assertEqual(ad_contracts[3].sid, 5)
797
798
        # Check that when the expiration_date has passed but the
799
        # notice_date hasn't, contract is still considered invalid.
800
        ad_contracts = finder.lookup_future_chain('AD', dt_3)
801
        self.assertEqual(len(ad_contracts), 1)
802
        self.assertEqual(ad_contracts[0].sid, 5)
803
804
        # Check that pd.NaT for as_of_date gives the whole chain
805
        ad_contracts = finder.lookup_future_chain('AD', pd.NaT)
806
        self.assertEqual(len(ad_contracts), 6)
807
        self.assertEqual(ad_contracts[5].sid, 5)
808
809
    def test_map_identifier_index_to_sids(self):
810
        # Build an empty finder and some Assets
811
        dt = pd.Timestamp('2014-01-01', tz='UTC')
812
        finder = self.asset_finder_type(self.env.engine)
813
        asset1 = Equity(1, symbol="AAPL")
814
        asset2 = Equity(2, symbol="GOOG")
815
        asset200 = Future(200, symbol="CLK15")
816
        asset201 = Future(201, symbol="CLM15")
817
818
        # Check for correct mapping and types
819
        pre_map = [asset1, asset2, asset200, asset201]
820
        post_map = finder.map_identifier_index_to_sids(pre_map, dt)
821
        self.assertListEqual([1, 2, 200, 201], post_map)
822
        for sid in post_map:
823
            self.assertIsInstance(sid, int)
824
825
        # Change order and check mapping again
826
        pre_map = [asset201, asset2, asset200, asset1]
827
        post_map = finder.map_identifier_index_to_sids(pre_map, dt)
828
        self.assertListEqual([201, 2, 200, 1], post_map)
829
830
    def test_compute_lifetimes(self):
831
        num_assets = 4
832
        trading_day = self.env.trading_day
833
        first_start = pd.Timestamp('2015-04-01', tz='UTC')
834
835
        frame = make_rotating_equity_info(
836
            num_assets=num_assets,
837
            first_start=first_start,
838
            frequency=self.env.trading_day,
839
            periods_between_starts=3,
840
            asset_lifetime=5
841
        )
842
843
        self.env.write_data(equities_df=frame)
844
        finder = self.env.asset_finder
845
846
        all_dates = pd.date_range(
847
            start=first_start,
848
            end=frame.end_date.max(),
849
            freq=trading_day,
850
        )
851
852
        for dates in all_subindices(all_dates):
853
            expected_with_start_raw = full(
854
                shape=(len(dates), num_assets),
855
                fill_value=False,
856
                dtype=bool,
857
            )
858
            expected_no_start_raw = full(
859
                shape=(len(dates), num_assets),
860
                fill_value=False,
861
                dtype=bool,
862
            )
863
864
            for i, date in enumerate(dates):
865
                it = frame[['start_date', 'end_date']].itertuples()
866
                for j, start, end in it:
867
                    # This way of doing the checks is redundant, but very
868
                    # clear.
869
                    if start <= date <= end:
870
                        expected_with_start_raw[i, j] = True
871
                        if start < date:
872
                            expected_no_start_raw[i, j] = True
873
874
            expected_with_start = pd.DataFrame(
875
                data=expected_with_start_raw,
876
                index=dates,
877
                columns=frame.index.values,
878
            )
879
            result = finder.lifetimes(dates, include_start_date=True)
880
            assert_frame_equal(result, expected_with_start)
881
882
            expected_no_start = pd.DataFrame(
883
                data=expected_no_start_raw,
884
                index=dates,
885
                columns=frame.index.values,
886
            )
887
            result = finder.lifetimes(dates, include_start_date=False)
888
            assert_frame_equal(result, expected_no_start)
889
890
    def test_sids(self):
891
        # Ensure that the sids property of the AssetFinder is functioning
892
        self.env.write_data(equities_identifiers=[1, 2, 3])
893
        sids = self.env.asset_finder.sids
894
        self.assertEqual(3, len(sids))
895
        self.assertTrue(1 in sids)
896
        self.assertTrue(2 in sids)
897
        self.assertTrue(3 in sids)
898
899
    def test_group_by_type(self):
900
        equities = make_simple_equity_info(
901
            range(5),
902
            start_date=pd.Timestamp('2014-01-01'),
903
            end_date=pd.Timestamp('2015-01-01'),
904
        )
905
        futures = make_commodity_future_info(
906
            first_sid=6,
907
            root_symbols=['CL'],
908
            years=[2014],
909
        )
910
        # Intersecting sid queries, to exercise loading of partially-cached
911
        # results.
912
        queries = [
913
            ([0, 1, 3], [6, 7]),
914
            ([0, 2, 3], [7, 10]),
915
            (list(equities.index), list(futures.index)),
916
        ]
917
        with tmp_asset_finder(equities=equities, futures=futures) as finder:
918
            for equity_sids, future_sids in queries:
919
                results = finder.group_by_type(equity_sids + future_sids)
920
                self.assertEqual(
921
                    results,
922
                    {'equity': set(equity_sids), 'future': set(future_sids)},
923
                )
924
925
    @parameterized.expand([
926
        (Equity, 'retrieve_equities', EquitiesNotFound),
927
        (Future, 'retrieve_futures_contracts', FutureContractsNotFound),
928
    ])
929
    def test_retrieve_specific_type(self, type_, lookup_name, failure_type):
930
        equities = make_simple_equity_info(
931
            range(5),
932
            start_date=pd.Timestamp('2014-01-01'),
933
            end_date=pd.Timestamp('2015-01-01'),
934
        )
935
        max_equity = equities.index.max()
936
        futures = make_commodity_future_info(
937
            first_sid=max_equity + 1,
938
            root_symbols=['CL'],
939
            years=[2014],
940
        )
941
        equity_sids = [0, 1]
942
        future_sids = [max_equity + 1, max_equity + 2, max_equity + 3]
943
        if type_ == Equity:
944
            success_sids = equity_sids
945
            fail_sids = future_sids
946
        else:
947
            fail_sids = equity_sids
948
            success_sids = future_sids
949
950
        with tmp_asset_finder(equities=equities, futures=futures) as finder:
951
            # Run twice to exercise caching.
952
            lookup = getattr(finder, lookup_name)
953
            for _ in range(2):
954
                results = lookup(success_sids)
955
                self.assertIsInstance(results, dict)
956
                self.assertEqual(set(results.keys()), set(success_sids))
957
                self.assertEqual(
958
                    valmap(int, results),
959
                    dict(zip(success_sids, success_sids)),
960
                )
961
                self.assertEqual(
962
                    {type_},
963
                    {type(asset) for asset in itervalues(results)},
964
                )
965
                with self.assertRaises(failure_type):
966
                    lookup(fail_sids)
967
                with self.assertRaises(failure_type):
968
                    # Should fail if **any** of the assets are bad.
969
                    lookup([success_sids[0], fail_sids[0]])
970
971
    def test_retrieve_all(self):
972
        equities = make_simple_equity_info(
973
            range(5),
974
            start_date=pd.Timestamp('2014-01-01'),
975
            end_date=pd.Timestamp('2015-01-01'),
976
        )
977
        max_equity = equities.index.max()
978
        futures = make_commodity_future_info(
979
            first_sid=max_equity + 1,
980
            root_symbols=['CL'],
981
            years=[2014],
982
        )
983
984
        with tmp_asset_finder(equities=equities, futures=futures) as finder:
985
            all_sids = finder.sids
986
            self.assertEqual(len(all_sids), len(equities) + len(futures))
987
            queries = [
988
                # Empty Query.
989
                (),
990
                # Only Equities.
991
                tuple(equities.index[:2]),
992
                # Only Futures.
993
                tuple(futures.index[:3]),
994
                # Mixed, all cache misses.
995
                tuple(equities.index[2:]) + tuple(futures.index[3:]),
996
                # Mixed, all cache hits.
997
                tuple(equities.index[2:]) + tuple(futures.index[3:]),
998
                # Everything.
999
                all_sids,
1000
                all_sids,
1001
            ]
1002
            for sids in queries:
1003
                equity_sids = [i for i in sids if i <= max_equity]
1004
                future_sids = [i for i in sids if i > max_equity]
1005
                results = finder.retrieve_all(sids)
1006
                self.assertEqual(sids, tuple(map(int, results)))
1007
1008
                self.assertEqual(
1009
                    [Equity for _ in equity_sids] +
1010
                    [Future for _ in future_sids],
1011
                    list(map(type, results)),
1012
                )
1013
                self.assertEqual(
1014
                    (
1015
                        list(equities.symbol.loc[equity_sids]) +
1016
                        list(futures.symbol.loc[future_sids])
1017
                    ),
1018
                    list(asset.symbol for asset in results),
1019
                )
1020
1021
    @parameterized.expand([
1022
        (EquitiesNotFound, 'equity', 'equities'),
1023
        (FutureContractsNotFound, 'future contract', 'future contracts'),
1024
        (SidsNotFound, 'asset', 'assets'),
1025
    ])
1026
    def test_error_message_plurality(self,
1027
                                     error_type,
1028
                                     singular,
1029
                                     plural):
1030
        try:
1031
            raise error_type(sids=[1])
1032
        except error_type as e:
1033
            self.assertEqual(
1034
                str(e),
1035
                "No {singular} found for sid: 1.".format(singular=singular)
1036
            )
1037
        try:
1038
            raise error_type(sids=[1, 2])
1039
        except error_type as e:
1040
            self.assertEqual(
1041
                str(e),
1042
                "No {plural} found for sids: [1, 2].".format(plural=plural)
1043
            )
1044
1045
1046
class AssetFinderCachedEquitiesTestCase(AssetFinderTestCase):
1047
1048
    def setUp(self):
1049
        self.env = TradingEnvironment(load=noop_load)
1050
        self.asset_finder_type = AssetFinderCachedEquities
1051
1052
1053
class TestFutureChain(TestCase):
1054
1055
    @classmethod
1056
    def setUpClass(cls):
1057
        metadata = {
1058
            0: {
1059
                'symbol': 'CLG06',
1060
                'root_symbol': 'CL',
1061
                'start_date': pd.Timestamp('2005-12-01', tz='UTC'),
1062
                'notice_date': pd.Timestamp('2005-12-20', tz='UTC'),
1063
                'expiration_date': pd.Timestamp('2006-01-20', tz='UTC')},
1064
            1: {
1065
                'root_symbol': 'CL',
1066
                'symbol': 'CLK06',
1067
                'start_date': pd.Timestamp('2005-12-01', tz='UTC'),
1068
                'notice_date': pd.Timestamp('2006-03-20', tz='UTC'),
1069
                'expiration_date': pd.Timestamp('2006-04-20', tz='UTC')},
1070
            2: {
1071
                'symbol': 'CLQ06',
1072
                'root_symbol': 'CL',
1073
                'start_date': pd.Timestamp('2005-12-01', tz='UTC'),
1074
                'notice_date': pd.Timestamp('2006-06-20', tz='UTC'),
1075
                'expiration_date': pd.Timestamp('2006-07-20', tz='UTC')},
1076
            3: {
1077
                'symbol': 'CLX06',
1078
                'root_symbol': 'CL',
1079
                'start_date': pd.Timestamp('2006-02-01', tz='UTC'),
1080
                'notice_date': pd.Timestamp('2006-09-20', tz='UTC'),
1081
                'expiration_date': pd.Timestamp('2006-10-20', tz='UTC')}
1082
        }
1083
1084
        env = TradingEnvironment(load=noop_load)
1085
        env.write_data(futures_data=metadata)
1086
        cls.asset_finder = env.asset_finder
1087
1088
    @classmethod
1089
    def tearDownClass(cls):
1090
        del cls.asset_finder
1091
1092
    def test_len(self):
1093
        """ Test the __len__ method of FutureChain.
1094
        """
1095
        # Sids 0, 1, & 2 have started, 3 has not yet started, but all are in
1096
        # the chain
1097
        cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL')
1098
        self.assertEqual(len(cl), 4)
1099
1100
        # Sid 0 is still valid on its notice date.
1101
        cl = FutureChain(self.asset_finder, lambda: '2005-12-20', 'CL')
1102
        self.assertEqual(len(cl), 4)
1103
1104
        # Sid 0 is now invalid, leaving Sids 1 & 2 valid (and 3 not started).
1105
        cl = FutureChain(self.asset_finder, lambda: '2005-12-21', 'CL')
1106
        self.assertEqual(len(cl), 3)
1107
1108
        # Sid 3 has started, so 1, 2, & 3 are now valid.
1109
        cl = FutureChain(self.asset_finder, lambda: '2006-02-01', 'CL')
1110
        self.assertEqual(len(cl), 3)
1111
1112
        # All contracts are no longer valid.
1113
        cl = FutureChain(self.asset_finder, lambda: '2006-09-21', 'CL')
1114
        self.assertEqual(len(cl), 0)
1115
1116
    def test_getitem(self):
1117
        """ Test the __getitem__ method of FutureChain.
1118
        """
1119
        cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL')
1120
        self.assertEqual(cl[0], 0)
1121
        self.assertEqual(cl[1], 1)
1122
        self.assertEqual(cl[2], 2)
1123
1124
        cl = FutureChain(self.asset_finder, lambda: '2005-12-20', 'CL')
1125
        self.assertEqual(cl[0], 0)
1126
1127
        cl = FutureChain(self.asset_finder, lambda: '2005-12-21', 'CL')
1128
        self.assertEqual(cl[0], 1)
1129
1130
        cl = FutureChain(self.asset_finder, lambda: '2006-02-01', 'CL')
1131
        self.assertEqual(cl[-1], 3)
1132
1133
    def test_iter(self):
1134
        """ Test the __iter__ method of FutureChain.
1135
        """
1136
        cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL')
1137
        for i, contract in enumerate(cl):
1138
            self.assertEqual(contract, i)
1139
1140
        # First contract is now invalid, so sids will be offset by one
1141
        cl = FutureChain(self.asset_finder, lambda: '2005-12-21', 'CL')
1142
        for i, contract in enumerate(cl):
1143
            self.assertEqual(contract, i + 1)
1144
1145
    def test_root_symbols(self):
1146
        """ Test that different variations on root symbols are handled
1147
        as expected.
1148
        """
1149
        # Make sure this successfully gets the chain for CL.
1150
        cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL')
1151
        self.assertEqual(cl.root_symbol, 'CL')
1152
1153
        # These root symbols don't exist, so RootSymbolNotFound should
1154
        # be raised immediately.
1155
        with self.assertRaises(RootSymbolNotFound):
1156
            FutureChain(self.asset_finder, lambda: '2005-12-01', 'CLZ')
1157
1158
        with self.assertRaises(RootSymbolNotFound):
1159
            FutureChain(self.asset_finder, lambda: '2005-12-01', '')
1160
1161
    def test_repr(self):
1162
        """ Test the __repr__ method of FutureChain.
1163
        """
1164
        cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL')
1165
        cl_feb = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL',
1166
                             as_of_date=pd.Timestamp('2006-02-01', tz='UTC'))
1167
1168
        # The default chain should not include the as of date.
1169
        self.assertEqual(repr(cl), "FutureChain(root_symbol='CL')")
1170
1171
        # An explicit as of date should show up in the repr.
1172
        self.assertEqual(
1173
            repr(cl_feb),
1174
            ("FutureChain(root_symbol='CL', "
1175
             "as_of_date='2006-02-01 00:00:00+00:00')")
1176
        )
1177
1178
    def test_as_of(self):
1179
        """ Test the as_of method of FutureChain.
1180
        """
1181
        cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL')
1182
1183
        # Test that the as_of_date is set correctly to the future
1184
        feb = pd.Timestamp('2006-02-01', tz='UTC')
1185
        cl_feb = cl.as_of(feb)
1186
        self.assertEqual(
1187
            cl_feb.as_of_date,
1188
            pd.Timestamp(feb, tz='UTC')
1189
        )
1190
1191
        # Test that the as_of_date is set correctly to the past, with
1192
        # args of str, datetime.datetime, and pd.Timestamp.
1193
        feb_prev = pd.Timestamp('2005-02-01', tz='UTC')
1194
        cl_feb_prev = cl.as_of(feb_prev)
1195
        self.assertEqual(
1196
            cl_feb_prev.as_of_date,
1197
            pd.Timestamp(feb_prev, tz='UTC')
1198
        )
1199
1200
        feb_prev = pd.Timestamp(datetime(year=2005, month=2, day=1), tz='UTC')
1201
        cl_feb_prev = cl.as_of(feb_prev)
1202
        self.assertEqual(
1203
            cl_feb_prev.as_of_date,
1204
            pd.Timestamp(feb_prev, tz='UTC')
1205
        )
1206
1207
        feb_prev = pd.Timestamp('2005-02-01', tz='UTC')
1208
        cl_feb_prev = cl.as_of(feb_prev)
1209
        self.assertEqual(
1210
            cl_feb_prev.as_of_date,
1211
            pd.Timestamp(feb_prev, tz='UTC')
1212
        )
1213
1214
        # Test that the as_of() method works with str args
1215
        feb_str = '2006-02-01'
1216
        cl_feb = cl.as_of(feb_str)
1217
        self.assertEqual(
1218
            cl_feb.as_of_date,
1219
            pd.Timestamp(feb, tz='UTC')
1220
        )
1221
1222
        # The chain as of the current dt should always be the same as
1223
        # the defualt chain.
1224
        self.assertEqual(cl[0], cl.as_of(pd.Timestamp('2005-12-01'))[0])
1225
1226
    def test_offset(self):
1227
        """ Test the offset method of FutureChain.
1228
        """
1229
        cl = FutureChain(self.asset_finder, lambda: '2005-12-01', 'CL')
1230
1231
        # Test that an offset forward sets as_of_date as expected
1232
        self.assertEqual(
1233
            cl.offset('3 days').as_of_date,
1234
            cl.as_of_date + pd.Timedelta(days=3)
1235
        )
1236
1237
        # Test that an offset backward sets as_of_date as expected, with
1238
        # time delta given as str, datetime.timedelta, and pd.Timedelta.
1239
        self.assertEqual(
1240
            cl.offset('-1000 days').as_of_date,
1241
            cl.as_of_date + pd.Timedelta(days=-1000)
1242
        )
1243
        self.assertEqual(
1244
            cl.offset(timedelta(days=-1000)).as_of_date,
1245
            cl.as_of_date + pd.Timedelta(days=-1000)
1246
        )
1247
        self.assertEqual(
1248
            cl.offset(pd.Timedelta('-1000 days')).as_of_date,
1249
            cl.as_of_date + pd.Timedelta(days=-1000)
1250
        )
1251
1252
        # An offset of zero should give the original chain.
1253
        self.assertEqual(cl[0], cl.offset(0)[0])
1254
        self.assertEqual(cl[0], cl.offset("0 days")[0])
1255
1256
        # A string that doesn't represent a time delta should raise a
1257
        # ValueError.
1258
        with self.assertRaises(ValueError):
1259
            cl.offset("blah")
1260
1261
    def test_cme_code_to_month(self):
1262
        codes = {
1263
            'F': 1,   # January
1264
            'G': 2,   # February
1265
            'H': 3,   # March
1266
            'J': 4,   # April
1267
            'K': 5,   # May
1268
            'M': 6,   # June
1269
            'N': 7,   # July
1270
            'Q': 8,   # August
1271
            'U': 9,   # September
1272
            'V': 10,  # October
1273
            'X': 11,  # November
1274
            'Z': 12   # December
1275
        }
1276
        for key in codes:
1277
            self.assertEqual(codes[key], cme_code_to_month(key))
1278
1279
    def test_month_to_cme_code(self):
1280
        codes = {
1281
            1: 'F',   # January
1282
            2: 'G',   # February
1283
            3: 'H',   # March
1284
            4: 'J',   # April
1285
            5: 'K',   # May
1286
            6: 'M',   # June
1287
            7: 'N',   # July
1288
            8: 'Q',   # August
1289
            9: 'U',   # September
1290
            10: 'V',  # October
1291
            11: 'X',  # November
1292
            12: 'Z',  # December
1293
        }
1294
        for key in codes:
1295
            self.assertEqual(codes[key], month_to_cme_code(key))
1296
1297
1298
class TestAssetDBVersioning(TestCase):
1299
1300
    def test_check_version(self):
1301
        env = TradingEnvironment(load=noop_load)
1302
        version_table = env.asset_finder.version_info
1303
1304
        # This should not raise an error
1305
        check_version_info(version_table, ASSET_DB_VERSION)
1306
1307
        # This should fail because the version is too low
1308
        with self.assertRaises(AssetDBVersionError):
1309
            check_version_info(version_table, ASSET_DB_VERSION - 1)
1310
1311
        # This should fail because the version is too high
1312
        with self.assertRaises(AssetDBVersionError):
1313
            check_version_info(version_table, ASSET_DB_VERSION + 1)
1314
1315
    def test_write_version(self):
1316
        env = TradingEnvironment(load=noop_load)
1317
        metadata = sa.MetaData(bind=env.engine)
1318
        version_table = _version_table_schema(metadata)
1319
        version_table.delete().execute()
1320
1321
        # Assert that the version is not present in the table
1322
        self.assertIsNone(sa.select((version_table.c.version,)).scalar())
1323
1324
        # This should fail because the table has no version info and is,
1325
        # therefore, consdered v0
1326
        with self.assertRaises(AssetDBVersionError):
1327
            check_version_info(version_table, -2)
1328
1329
        # This should not raise an error because the version has been written
1330
        write_version_info(version_table, -2)
1331
        check_version_info(version_table, -2)
1332
1333
        # Assert that the version is in the table and correct
1334
        self.assertEqual(sa.select((version_table.c.version,)).scalar(), -2)
1335
1336
        # Assert that trying to overwrite the version fails
1337
        with self.assertRaises(sa.exc.IntegrityError):
1338
            write_version_info(version_table, -3)
1339
1340
    def test_finder_checks_version(self):
1341
        # Create an env and give it a bogus version number
1342
        env = TradingEnvironment(load=noop_load)
1343
        metadata = sa.MetaData(bind=env.engine)
1344
        version_table = _version_table_schema(metadata)
1345
        version_table.delete().execute()
1346
        write_version_info(version_table, -2)
1347
        check_version_info(version_table, -2)
1348
1349
        # Assert that trying to build a finder with a bad db raises an error
1350
        with self.assertRaises(AssetDBVersionError):
1351
            AssetFinder(engine=env.engine)
1352
1353
        # Change the version number of the db to the correct version
1354
        version_table.delete().execute()
1355
        write_version_info(version_table, ASSET_DB_VERSION)
1356
        check_version_info(version_table, ASSET_DB_VERSION)
1357
1358
        # Now that the versions match, this Finder should succeed
1359
        AssetFinder(engine=env.engine)
1360