def test_single_source(self): # Just using the built-in defaults. See # zipline.sources.py source = SpecificEquityTrades() expected = list(source) source.rewind() # The raw source doesn't handle done messaging, so we need to # append a done message for sort to work properly. with_done = chain(source, [done_message(source.get_hash())]) self.run_date_sort(with_done, expected, [source.get_hash()])
def test_multi_source(self): filter = [2, 3] args_a = tuple() kwargs_a = { 'count': 100, 'sids': [1, 2, 3], 'start': datetime(2012, 1, 3, 15, tzinfo=pytz.utc), 'delta': timedelta(minutes=6), 'filter': filter } source_a = SpecificEquityTrades(*args_a, **kwargs_a) args_b = tuple() kwargs_b = { 'count': 100, 'sids': [2, 3, 4], 'start': datetime(2012, 1, 3, 15, tzinfo=pytz.utc), 'delta': timedelta(minutes=5), 'filter': filter } source_b = SpecificEquityTrades(*args_b, **kwargs_b) all_events = list(chain(source_a, source_b)) # The expected output is all events, sorted by dt with # source_id as a tiebreaker. expected = sorted(all_events, comp) source_ids = [source_a.get_hash(), source_b.get_hash()] # Generating the events list consumes the sources. Rewind them # for testing. source_a.rewind() source_b.rewind() # Append a done message to each source. with_done_a = chain(source_a, [done_message(source_a.get_hash())]) with_done_b = chain(source_b, [done_message(source_b.get_hash())]) interleaved = alternate(with_done_a, with_done_b) # Test sort with alternating messages from source_a and # source_b. self.run_date_sort(interleaved, expected, source_ids) source_a.rewind() source_b.rewind() with_done_a = chain(source_a, [done_message(source_a.get_hash())]) with_done_b = chain(source_b, [done_message(source_b.get_hash())]) sequential = chain(with_done_a, with_done_b) # Test sort with all messages from a, followed by all messages # from b. self.run_date_sort(sequential, expected, source_ids)
def test_sort_composite(self): filter = [1, 2] #Set up source a. One hour between events. args_a = tuple() kwargs_a = { 'count': 100, 'sids': [1], 'start': datetime(2012, 6, 6, 0), 'delta': timedelta(hours=1), 'filter': filter } source_a = SpecificEquityTrades(*args_a, **kwargs_a) #Set up source b. One day between events. args_b = tuple() kwargs_b = { 'count': 50, 'sids': [2], 'start': datetime(2012, 6, 6, 0), 'delta': timedelta(days=1), 'filter': filter } source_b = SpecificEquityTrades(*args_b, **kwargs_b) #Set up source c. One minute between events. args_c = tuple() kwargs_c = { 'count': 150, 'sids': [1, 2], 'start': datetime(2012, 6, 6, 0), 'delta': timedelta(minutes=1), 'filter': filter } source_c = SpecificEquityTrades(*args_c, **kwargs_c) # Set up source d. This should produce no events because the # internal sids don't match the filter. args_d = tuple() kwargs_d = { 'count': 50, 'sids': [3], 'start': datetime(2012, 6, 6, 0), 'delta': timedelta(minutes=1), 'filter': filter } source_d = SpecificEquityTrades(*args_d, **kwargs_d) sources = [source_a, source_b, source_c, source_d] hashes = [source.get_hash() for source in sources] sort_out = date_sorted_sources(*sources) # Read all the values from sort and assert that they arrive in # the correct sorting with the expected hash values. to_list = list(sort_out) copy = to_list[:] # We should have 300 events (100 from a, 150 from b, 50 from c) assert len(to_list) == 300 for e in to_list: # All events should match one of our expected source_ids. assert e.source_id in hashes # But none of them should match source_d. assert e.source_id != source_d.get_hash() # The events should be sorted by dt, with source_id as tiebreaker. expected = sorted(copy, comp) assert to_list == expected