mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-28 21:10:25 +08:00
fawce
162 lines
4.1 KiB
Python
162 lines
4.1 KiB
Python
import zmq
|
|
import pytz
|
|
from datetime import datetime, timedelta
|
|
|
|
from unittest2 import TestCase
|
|
from collections import defaultdict
|
|
from zipline.gens.composites import date_sorted_sources
|
|
|
|
from zipline.finance.trading import SIMULATION_STYLE
|
|
from zipline.core.devsimulator import AddressAllocator
|
|
|
|
from zipline.utils.test_utils import (
|
|
setup_logger,
|
|
teardown_logger,
|
|
create_monitor,
|
|
launch_monitor
|
|
)
|
|
|
|
from zipline.core import Component
|
|
from zipline.protocol import (
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME
|
|
)
|
|
|
|
from zipline.gens.tradegens import SpecificEquityTrades
|
|
|
|
import logbook
|
|
log = logbook.Logger('ComponentTestCase')
|
|
|
|
allocator = AddressAllocator(1000)
|
|
|
|
|
|
class ComponentTestCase(TestCase):
|
|
|
|
leased_sockets = defaultdict(list)
|
|
|
|
def setUp(self):
|
|
self.zipline_test_config = {
|
|
'allocator' : allocator,
|
|
'sid' : 133,
|
|
'devel' : False,
|
|
'results_socket' : allocator.lease(1)[0],
|
|
'simulation_style' : SIMULATION_STYLE.FIXED_SLIPPAGE
|
|
}
|
|
self.ctx = zmq.Context()
|
|
setup_logger(self)
|
|
|
|
def tearDown(self):
|
|
teardown_logger(self)
|
|
|
|
def test_source(self):
|
|
monitor = create_monitor(allocator)
|
|
socket_uri = allocator.lease(1)[0]
|
|
count = 100
|
|
|
|
filter = [1,2,3,4]
|
|
#Set up source a. One minute between events.
|
|
args_a = tuple()
|
|
kwargs_a = {
|
|
'sids' : [1,2],
|
|
'start' : datetime(2012,6,6,0,tzinfo=pytz.utc),
|
|
'delta' : timedelta(minutes = 1),
|
|
'filter' : filter,
|
|
'count' : count
|
|
}
|
|
|
|
comp_a = Component(
|
|
SpecificEquityTrades(*args_a, **kwargs_a),
|
|
monitor,
|
|
socket_uri,
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME
|
|
)
|
|
|
|
launch_monitor(monitor)
|
|
iter_a = iter(comp_a)
|
|
ev = iter_a.next()
|
|
return
|
|
|
|
for event in comp_a:
|
|
log.info(event)
|
|
|
|
|
|
def test_sort(self):
|
|
monitor = create_monitor(allocator)
|
|
socket_uris = allocator.lease(3)
|
|
count = 100
|
|
|
|
filter = [1,2,3,4]
|
|
#Set up source a. One minute between events.
|
|
args_a = tuple()
|
|
kwargs_a = {
|
|
'sids' : [1,2],
|
|
'start' : datetime(2012,6,6,0,tzinfo=pytz.utc),
|
|
'delta' : timedelta(minutes = 1),
|
|
'filter' : filter,
|
|
'count' : count
|
|
}
|
|
|
|
|
|
comp_a = Component(
|
|
SpecificEquityTrades(*args_a, **kwargs_a),
|
|
monitor,
|
|
socket_uris[0],
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME
|
|
)
|
|
|
|
|
|
#Set up source b. Two minutes between events.
|
|
args_b = tuple()
|
|
kwargs_b = {
|
|
'sids' : [2],
|
|
'start' : datetime(2012,1,3,15, tzinfo = pytz.utc),
|
|
'delta' : timedelta(minutes = 1),
|
|
'filter' : filter,
|
|
'count' : count
|
|
}
|
|
|
|
|
|
comp_b = Component(
|
|
SpecificEquityTrades(*args_b, **kwargs_b),
|
|
monitor,
|
|
socket_uris[1],
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME
|
|
)
|
|
|
|
#Set up source c. Three minutes between events.
|
|
args_c = tuple()
|
|
kwargs_c = {
|
|
'sids' : [3],
|
|
'start' : datetime(2012,1,3,15, tzinfo = pytz.utc),
|
|
'delta' : timedelta(minutes = 1),
|
|
'filter' : filter,
|
|
'count' : count
|
|
}
|
|
|
|
comp_c = Component(
|
|
SpecificEquityTrades(*args_c, **kwargs_c),
|
|
monitor,
|
|
socket_uris[2],
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME
|
|
)
|
|
|
|
launch_monitor(monitor)
|
|
sources = [comp_a, comp_b, comp_c]
|
|
gens = [iter(source) for source in sources]
|
|
sorted_out = date_sorted_sources(gens)
|
|
|
|
prev = None
|
|
sort_count = 0
|
|
for msg in sorted_out:
|
|
if prev:
|
|
self.assertTrue(msg.dt >= prev.dt, \
|
|
"Messages should be in date ascending order")
|
|
prev = msg
|
|
sort_count += 1
|
|
|
|
self.assertEqual(count*3, sort_count)
|