mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-29 01:57:31 +08:00
fawce
107 lines
2.8 KiB
Python
107 lines
2.8 KiB
Python
import zmq
|
|
import pytz
|
|
from datetime import datetime, timedelta
|
|
|
|
from unittest2 import TestCase
|
|
from collections import defaultdict
|
|
|
|
from zipline.finance.trading import SIMULATION_STYLE
|
|
from zipline.core.devsimulator import AddressAllocator
|
|
from zipline.lines import SimulatedTrading
|
|
|
|
from zipline.utils.test_utils import (
|
|
drain_zipline,
|
|
check,
|
|
setup_logger,
|
|
teardown_logger,
|
|
launch_component,
|
|
create_monitor,
|
|
launch_monitor
|
|
)
|
|
|
|
|
|
from zipline.core import Component
|
|
from zipline.core.component import ComponentSocketArgs
|
|
from zipline.protocol import (
|
|
DATASOURCE_FRAME,
|
|
DATASOURCE_UNFRAME
|
|
)
|
|
|
|
from zipline.gens.tradegens import SpecificEquityTrades
|
|
from zipline.gens.utils import hash_args
|
|
from zipline.gens.zmqgen import gen_from_poller
|
|
|
|
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):
|
|
self.ctx.term()
|
|
teardown_logger(self)
|
|
|
|
def test_specific_equity_source(self):
|
|
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' : 100
|
|
}
|
|
|
|
c_id = SpecificEquityTrades.__name__ + hash_args(args_a, kwargs_a)
|
|
mon = create_monitor(allocator)
|
|
|
|
out_socket_args = ComponentSocketArgs(
|
|
style=zmq.PUSH,
|
|
uri=allocator.lease(1)[0],
|
|
bind=True
|
|
)
|
|
|
|
c = Component(
|
|
SpecificEquityTrades,
|
|
args_a,
|
|
kwargs_a,
|
|
c_id,
|
|
out_socket_args,
|
|
DATASOURCE_FRAME,
|
|
mon
|
|
)
|
|
|
|
mon.manage(set([c.get_id]))
|
|
mon_proc = launch_monitor(mon)
|
|
|
|
# launch in a process
|
|
proc = launch_component(c)
|
|
|
|
pull_socket = self.ctx.socket(zmq.PULL)
|
|
pull_socket.connect(out_socket_args.uri)
|
|
poller = zmq.Poller()
|
|
poller.register(pull_socket, zmq.POLLIN)
|
|
unframe = DATASOURCE_UNFRAME
|
|
for msg in gen_from_poller(poller, pull_socket, unframe):
|
|
# assert things about the messages.
|
|
log.info(msg)
|
|
|
|
pull_socket.close()
|
|
log.info("DONE!")
|