diff --git a/zipline/core.py b/zipline/core.py deleted file mode 100644 index 7cd87030..00000000 --- a/zipline/core.py +++ /dev/null @@ -1,56 +0,0 @@ -""" -Provides simulated data feed services... -""" -import multiprocessing -from gevent_zeromq import zmq -import json -import copy -import threading -import datetime -import atexit - -import zipline.util as qutil -import zipline.messaging as qmsg - -class SimulatorBase(ComponentHost): - """ - Simulator coordinates the launch and communication of source, feed, transform, and merge components. - """ - - def __init__(self, addresses): - """ - """ - self.feed = ParallelBuffer() - self.merge = MergedParallelBuffer() - - #workaround for defect in threaded use of strptime: http://bugs.python.org/issue11108 - qutil.parse_date("2012/02/13-10:04:28.114") - - #register the feed and the merge - self.register_components([self.feed, self.merge]) - - def simulate(self): - self.launch_component(self) - -class ThreadSimulator(SimulatorBase): - - def __init__(self, sources, transforms, client, feed=None, merge=None): - SimulatorBase.__init__(self, sources, transforms, client, feed, merge) - - def launch_component(self, name, component): - qutil.LOGGER.info("starting {name}".format(name=name)) - thread = threading.Thread(target=component.run) - thread.start() - return thread - -class ProcessSimulator(SimulatorBase): - - def __init__(self, sources, transforms, client, feed=None, merge=None): - SimulatorBase.__init__(self, sources, transforms, client, feed, merge) - - def launch_component(self, name, component): - qutil.LOGGER.info("starting {name}".format(name=name)) - proc = multiprocessing.Process(target=component.run) - proc.start() - return proc - diff --git a/zipline/messaging.py b/zipline/messaging.py index 1e3fd3b2..8c5076d1 100644 --- a/zipline/messaging.py +++ b/zipline/messaging.py @@ -3,7 +3,9 @@ Commonly used messaging components. """ import json import uuid -from gevent_zeromq import zmq +import datetime +#from gevent_zeromq import zmq +import zmq import zipline.util as qutil class Component(object): @@ -28,174 +30,223 @@ class Component(object): Bind/Connect methods will return the correct socket type for each address. Any sockets on which recv is expected to be called will also return a Poller. - """ - self.context = zmq.Context() + """ + self.zmq = None + self.context = None self.addresses = addresses - self.sockets = [] + self.out_socket = None def get_id(self): - raise NotImplemented + NotImplemented def open(self): - raise NotImplemented + NotImplemented def do_work(self): - raise NotImplemented + NotImplemented def run(self): try: + #TODO: can't initialize a boolean in the __init__? + self.done = False + #TODO: figure out why sockets can't be set in the __init__ + self.sockets = [] + #TODO: figure out why addresses can't be set in the __init__ + self.addresses = {'sync_address' : "tcp://127.0.0.1:{port}".format(port=10100), + 'data_address' : "tcp://127.0.0.1:{port}".format(port=10101), + 'feed_address' : "tcp://127.0.0.1:{port}".format(port=10102), + 'merge_address' : "tcp://127.0.0.1:{port}".format(port=10103), + 'result_address' : "tcp://127.0.0.1:{port}".format(port=10104) + } + + self.context = self.zmq.Context() self.open() - self.connect_sync() - while self.confirm(): - self.do_work() - #notify host we're done - self.sync_socket.send(self.sync_id + ":DONE") + self.setup_sync() + self.loop() #close all the sockets for sock in self.sockets: sock.close() finally: self.context.destroy() + + def loop(self): + while not self.done: + self.confirm() + self.do_work() + + def signal_done(self): + #notify down stream components that we're done + if(self.out_socket != None): + self.out_socket.send("DONE") + #notify host we're done + self.sync_socket.send(self.get_id() + ":DONE") + self.receive_sync_ack() + #notify internal work look that we're done + self.done = True + + def is_done_message(self, message): + return message == "DONE" + + def confirm(self): + # send a synchronization request to the host + self.sync_socket.send(self.get_id() + ":RUN") + self.receive_sync_ack() + + def receive_sync_ack(self): + # wait for synchronization reply from the host + socks = dict(self.sync_poller.poll(2000)) #timeout after 2 seconds. - def confirm(self): - try: - # send a synchronization request to the host - self.sync_socket.send(self.sync_id + ":RUNNING") - # wait for synchronization reply from the host - socks = dict(self.poller.poll(2000)) #timeout after 2 seconds. - - if self.sync_socket in socks and socks[self.sync_socket] == zmq.POLLIN: - message = self.sync_socket.recv() - - return True - except: - qutil.LOGGER.exception("exception in confirmation for {source}. Exiting.".format(source=self.sync_id)) - return False + if self.sync_socket in socks and socks[self.sync_socket] == self.zmq.POLLIN: + message = self.sync_socket.recv() def bind_data(self): - return self.bind_pull_socket(self, self.addresses['data_address']) + return self.bind_pull_socket(self.addresses['data_address']) def connect_data(self): - return self.connect_push_socket(self, self.addresses['data_address']) + return self.connect_push_socket(self.addresses['data_address']) def bind_feed(self): - return self.bind_pub_socket(self, self.addresses['feed_address']) + return self.bind_pub_socket(self.addresses['feed_address']) def connect_feed(self): - return self.bind_sub_socket(self, self.addresses['feed_address']) + return self.connect_sub_socket(self.addresses['feed_address']) def bind_merge(self): - return self.bind_pull_socket(self, self.addresses['merge_address']) + return self.bind_pull_socket(self.addresses['merge_address']) def connect_merge(self): - return self.connect_push_socket(self, self.addresses['merge_address']) + return self.connect_push_socket(self.addresses['merge_address']) def bind_result(self): - return self.bind_pub_socket(self, self.addresses['result_address']) + return self.bind_pub_socket(self.addresses['result_address']) def connect_result(self): - return self.bind_sub_socket(self, self.addresses['result_address']) + return self.connect_sub_socket(self.addresses['result_address']) def bind_pull_socket(self, addr): - pull_socket = self.context.socket(zmq.PULL) + pull_socket = self.context.socket(self.zmq.PULL) pull_socket.bind(addr) - poller = zmq.Poller() - poller.register(self.pull_socket, zmq.POLLIN) + poller = self.zmq.Poller() + poller.register(pull_socket, self.zmq.POLLIN) self.sockets.append(pull_socket) return pull_socket, poller def connect_push_socket(self, addr): - push_socket = self.context.socket(zmq.PUSH) - push_socket.connect(self.merge_address) - #push_socket.setsockopt(zmq.LINGER,0) + push_socket = self.context.socket(self.zmq.PUSH) + push_socket.connect(addr) + #push_socket.setsockopt(self.zmq.LINGER,0) self.sockets.append(push_socket) + self.out_socket = push_socket return push_socket def bind_pub_socket(self, addr): - pub_socket = self.context.socket(zmq.PUB) - pub_socket.bind(self.pub_address) - #pub_socket.setsockopt(zmq.LINGER,0) - poller = zmq.Poller() - poller.register(self.pub_socket, zmq.POLLIN) - self.sockets.append(pub_socket) - return pub_socket, poller + pub_socket = self.context.socket(self.zmq.PUB) + pub_socket.bind(addr) + #pub_socket.setsockopt(self.zmq.LINGER,0) + self.out_socket = pub_socket + return pub_socket def connect_sub_socket(self, addr): - sub_socket = self.context.socket(zmq.SUB) - sub_socket.connect(self.feed_address) - sub_socket.setsockopt(zmq.SUBSCRIBE,'') + sub_socket = self.context.socket(self.zmq.SUB) + sub_socket.connect(addr) + sub_socket.setsockopt(self.zmq.SUBSCRIBE,'') + poller = self.zmq.Poller() + poller.register(sub_socket, self.zmq.POLLIN) self.sockets.append(sub_socket) - return sub_socket - - def bind_sync(self): - sync_socket = self.context.socket(zmq.REP) - sync_socket.bind(self.addresses['sync_address']) - poller = zmq.Poller() - poller.register(self.sync_socket, zmq.POLLIN) - self.sockets.append(sync_socket) - return sync_socket, poller - - def connect_sync(self): - self.sync_socket = self.context.socket(zmq.REQ) + return sub_socket, poller + + def setup_sync(self): + qutil.LOGGER.debug("Connecting sync client for {id}".format(id=self.get_id())) + self.sync_socket = self.context.socket(self.zmq.REQ) self.sync_socket.connect(self.addresses['sync_address']) - self.sync_socket.setsockopt(zmq.LINGER,0) - self.poller = zmq.Poller() - self.poller.register(self.sync_socket, zmq.POLLIN) - self.sockets.append(sync_socket) + #self.sync_socket.setsockopt(self.zmq.LINGER,0) + self.sync_poller = self.zmq.Poller() + self.sync_poller.register(self.sync_socket, self.zmq.POLLIN) + self.sockets.append(self.sync_socket) class ComponentHost(Component): """Component that can launch multiple sub-components, synchronize their start, and then wait for all components to be finished.""" - def __init__(self, addresses): + def __init__(self, addresses, gevent_needed=False): Component.__init__(self, addresses) - self.components = {} - self.timeout = datetime.timedelta(seconds=5) + if gevent_needed: + module = __import__('gevent_zmq', 'zmq') + else: + module = __import__('zmq') + self.zmq = module + #workaround for defect in threaded use of strptime: http://bugs.python.org/issue11108 + qutil.parse_date("2012/02/13-10:04:28.114") + self.components = {} + self.sync_register = {} + self.timeout = datetime.timedelta(seconds=5) + self.feed = ParallelBuffer() + self.merge = MergedParallelBuffer() + self.passthrough = PassthroughTransform() + + #register the feed and the merge + self.register_components([self.feed, self.merge, self.passthrough]) def register_components(self, component_list): for component in component_list: + component.zmq = self.zmq self.components[component.get_id()] = component + self.sync_register[component.get_id()] = datetime.datetime.utcnow() + if(isinstance(component, DataSource)): + self.feed.add_source(component.get_id()) + if(isinstance(component, BaseTransform)): + self.merge.add_source(component.get_id()) def unregister_component(self, component_id): del(self.components[component_id]) + del(self.sync_register[component_id]) + + def setup_sync(self): + """Start the sync server.""" + qutil.LOGGER.debug("Connecting sync server.") + self.sync_socket = self.context.socket(self.zmq.REP) + self.sync_socket.bind(self.addresses['sync_address']) + self.poller = self.zmq.Poller() + self.poller.register(self.sync_socket, self.zmq.POLLIN) + self.sockets.append(self.sync_socket) def open(self): - self.sync_socket, self.poller = self.bind_sync() for component in self.components.values(): self.launch_component(component) def is_timed_out(self): cur_time = datetime.datetime.utcnow() - if(len(self.sync_register) == 0): + if(len(self.components) == 0): qutil.LOGGER.info("Component register is empty.") return True for source, last_dt in self.sync_register.iteritems(): if((cur_time - last_dt) > self.timeout): - qutil.LOGGER.info("Time out for {source}. Current registery: {reg}".format(source=source, reg=self.sync_register)) + qutil.LOGGER.info("Time out for {source}. Current component registery: {reg}".format(source=source, reg=self.components)) return True return False - def run(self): + def loop(self): while not self.is_timed_out(): # wait for synchronization request socks = dict(self.poller.poll(2000)) #timeout after 2 seconds. - if self.sync_socket in socks and socks[self.sync_socket] == zmq.POLLIN: - try: - msg = self.sync_socket.recv() - parts = msg.split(':') - sync_id = parts[0] - status = parts[1] - if(status == "DONE"): - self.unregister_component(sync_id) - else: - self.sync_register[sync_id] = datetime.datetime.utcnow() - #qutil.LOGGER.info("confirmed {id}".format(id=msg)) - # send synchronization reply - self.sync_socket.send('ack', zmq.NOBLOCK) - except: - qutil.LOGGER.exception("Exception in sync components loop") + if self.sync_socket in socks and socks[self.sync_socket] == self.zmq.POLLIN: + msg = self.sync_socket.recv() + parts = msg.split(':') + if(len(parts) < 2): + qutil.LOGGER.info("got bad confirm: {msg}".format(msg=msg)) + sync_id = parts[0] + status = parts[1] + if(self.is_done_message(status)): + qutil.LOGGER.info("{id} is DONE".format(id=sync_id)) + self.unregister_component(sync_id) + else: + self.sync_register[sync_id] = datetime.datetime.utcnow() + #qutil.LOGGER.info("confirmed {id}".format(id=msg)) + # send synchronization reply + self.sync_socket.send('ack', self.zmq.NOBLOCK) - def luanch_component(self, component): - raise NotImplemented + def launch_component(self, component): + NotImplemented class ParallelBuffer(Component): """Connects to N PULL sockets, publishing all messages received to a PUB socket. @@ -224,21 +275,17 @@ class ParallelBuffer(Component): # wait for synchronization reply from the host socks = dict(self.poller.poll(2000)) #timeout after 2 seconds. - if self.pull_socket in socks and socks[self.pull_socket] == zmq.POLLIN: + if self.pull_socket in socks and socks[self.pull_socket] == self.zmq.POLLIN: message = self.pull_socket.recv() - - event = json.loads(message) - if(event["type"] == "DONE"): - ds_finished_counter += 1 - if(len(self.data_buffer) == ds_finished_counter): + if(self.is_done_message(message)): + self.ds_finished_counter += 1 + if(len(self.data_buffer) == self.ds_finished_counter): #drain any remaining messages in the buffer self.drain() - #send the DONE message - self.feed_socket.send("DONE", zmq.NOBLOCK) - qutil.LOGGER.info("received {n} messages, sent {m} messages".format(n=self.received_count, - m=self.sent_count)) + self.signal_done() else: - self.append(event[u's'], event) + event = json.loads(message) + self.append(event[u'id'], event) self.send_next() def __len__(self): @@ -294,7 +341,7 @@ class ParallelBuffer(Component): event = self.next() if(event != None): - self.feed_socket.send(json.dumps(event), zmq.NOBLOCK) + self.feed_socket.send(json.dumps(event), self.zmq.NOBLOCK) self.sent_count += 1 @@ -305,21 +352,29 @@ class MergedParallelBuffer(ParallelBuffer): def __init__(self): ParallelBuffer.__init__(self) + + def open(self): + self.pull_socket, self.poller = self.bind_merge() + self.feed_socket = self.bind_result() def next(self): """Get the next merged message from the feed buffer.""" if(not(self.is_full() or self.draining)): return - result = self.feed.pop(0) + #get the raw event from the passthrough transform. + result = self.data_buffer["PASSTHROUGH"].pop(0)['value'] for source, events in self.data_buffer.iteritems(): - if(source == "feed"): + if(source == "PASSTHROUGH"): continue if(len(events) > 0): cur = events.pop(0) result[source] = cur['value'] return result + def get_id(self): + return "MERGE" + class BaseTransform(Component): """Top level execution entry point for the transform:: @@ -346,11 +401,9 @@ class BaseTransform(Component): """ #create the feed. self.feed_socket, self.poller = self.connect_feed() - #create the result PUSH self.result_socket = self.connect_merge() - def do_work(self): """ Loops until feed's DONE message is received: @@ -358,21 +411,17 @@ class BaseTransform(Component): - call transform (subclass' method) on event - send the transformed event """ - qutil.LOGGER.info("starting {name} event loop".format(name = self.state['name'])) socks = dict(self.poller.poll(2000)) #timeout after 2 seconds. - if self.feed_socket in socks and socks[self.feed_socket] == zmq.POLLIN: + if self.feed_socket in socks and socks[self.feed_socket] == self.zmq.POLLIN: message = self.feed_socket.recv() - if(message == "DONE"): - qutil.LOGGER.info("{name} received the Done message from the feed".format(name=self.state['name'])) - self.result_socket.send("DONE", zmq.NOBLOCK) + if(self.is_done_message(message)): + self.signal_done() return - self.received_count += 1 event = json.loads(message) cur_state = self.transform(event) cur_state['dt'] = event['dt'] - cur_state['name'] = self.state['name'] - self.result_socket.send(json.dumps(cur_state), zmq.NOBLOCK) - self.sent_count += 1 + cur_state['id'] = self.state['name'] + self.result_socket.send(json.dumps(cur_state), self.zmq.NOBLOCK) def transform(self, event): """ Must return the transformed value as a map with {name:"name of new transform", value: "value of new field"} @@ -381,5 +430,43 @@ class BaseTransform(Component): transformed value: self.state['value'] = transformed_value """ - return {'value'=event} - \ No newline at end of file + NotImplemented + +class PassthroughTransform(BaseTransform): + + def __init__(self): + BaseTransform.__init__(self, "PASSTHROUGH") + + def transform(self, event): + return {'value':event} + +class DataSource(Component): + """ + Baseclass for data sources. Subclass and implement send_all - usually this + means looping through all records in a store, converting to a dict, and + calling send(map). + """ + def __init__(self, source_id, addresses): + Component.__init__(self, addresses) + self.id = source_id + self.cur_event = None + + def get_id(self): + return self.id + + def open(self): + #create the data sink. Based on http://zguide.zeromq.org/py:tasksink2 + self.data_socket = self.connect_data() + + def send(self, event): + """ + event is expected to be a dict + sets id and type properties in the dict + sends to the data_socket. + """ + event['id'] = self.id + event['type'] = self.get_type() + self.data_socket.send(json.dumps(event)) + + def get_type(self): + raise NotImplemented \ No newline at end of file diff --git a/zipline/simulator.py b/zipline/simulator.py new file mode 100644 index 00000000..c64c9a5f --- /dev/null +++ b/zipline/simulator.py @@ -0,0 +1,49 @@ +""" +Provides simulated data feed services... +""" +import multiprocessing +import json +import copy +import threading + +import zipline.util as qutil +import zipline.messaging as qmsg + +class SimulatorBase(qmsg.ComponentHost): + """ + Simulator coordinates the launch and communication of source, feed, transform, and merge components. + """ + + def __init__(self, addresses, gevent_needed=False): + """ + """ + qmsg.ComponentHost.__init__(self, addresses, gevent_needed) + + def simulate(self): + self.run() + + def get_id(self): + return "Simulator" + +class ThreadSimulator(SimulatorBase): + + def __init__(self, addresses): + SimulatorBase.__init__(self, addresses) + + def launch_component(self, component): + qutil.LOGGER.info("starting {name}".format(name=component.get_id())) + thread = threading.Thread(target=component.run) + thread.start() + return thread + +class ProcessSimulator(SimulatorBase): + + def __init__(self, addresses): + SimulatorBase.__init__(self, addresses) + + def launch_component(self, component): + qutil.LOGGER.info("starting {name}".format(name=component.get_id())) + proc = multiprocessing.Process(target=component.run) + proc.start() + return proc + diff --git a/zipline/sources.py b/zipline/sources.py index 5666db00..164576cb 100644 --- a/zipline/sources.py +++ b/zipline/sources.py @@ -2,123 +2,38 @@ Provides data handlers that can push messages to a zipline.core.DataFeed """ import datetime -from gevent_zeromq import zmq import json import random import zipline.util as qutil -import zipline.messaging as qmsg - -class DataSource(qmsg.Component): - """ - Baseclass for data sources. Subclass and implement send_all - usually this - means looping through all records in a store, converting to a dict, and - calling send(map). - """ - def __init__(self, source_id): - self.id = source_id - self.host = host - self.data_address = None - self.sync = None - self.cur_event = None - self.context = None - self.data_socket = None - - def get_id(self): - return self.id - - def set_addresses(self, addresses): - self.data_address = addresses['data_address'] - - def open(self): - """create zmq context and socket""" - qutil.LOGGER.info( - "starting data source:{id} on {addr}" - .format(id=self.id, addr=self.data_address)) - - self.context = zmq.Context() - - #create the data sink. Based on http://zguide.zeromq.org/py:tasksink2 - self.data_socket = self.context.socket(zmq.PUSH) - self.data_socket.connect(self.data_address) - #self.data_socket.setsockopt(zmq.LINGER,0) - - self.sync.open() - - def run(self): - """Fully execute this datasource.""" - try: - self.open() - self.send_all() - except: - qutil.LOGGER.info("Exception running datasource.") - finally: - self.close() - - def send_all(self): - """Subclasses must implement this method.""" - raise NotImplementedError() - - def send(self, event): - """ - event is expected to be a dict - sets id and type properties in the dict - sends to the data_socket. - """ - event['id'] = self.id - event['type'] = self.get_type() - self.data_socket.send(json.dumps(event), zmq.NOBLOCK) - - def get_type(self): - raise NotImplemented - - def close(self): - """ - Close the zmq context and sockets. - """ - qutil.LOGGER.info("sending DONE message.") - try: - done_msg = {} - done_msg['type'] = 'DONE' - done_msg['s'] = self.id - self.data_socket.send(json.dumps(done_msg), zmq.NOBLOCK) - - qutil.LOGGER.info("closing data socket") - self.data_socket.close() - qutil.LOGGER.info("closing sync") - self.sync.close() - qutil.LOGGER.info("closing context") - except: - qutil.LOGGER.exception("failed to send DONE message") - finally: - self.context.destroy() - qutil.LOGGER.info("finished processing data source") +import zipline.messaging as qmsg -class RandomEquityTrades(DataSource): +class RandomEquityTrades(qmsg.DataSource): """Generates a random stream of trades for testing.""" - def __init__(self, sid, source_id, count): - DataSource.__init__(self, source_id) - self.count = count - self.sid = sid + def __init__(self, sid, source_id, count, addresses): + qmsg.DataSource.__init__(self, source_id, addresses) + self.count = count + self.incr = 0 + self.sid = sid + self.trade_start = datetime.datetime.now() + self.minute = datetime.timedelta(minutes=1) + self.price = random.uniform(5.0, 50.0) def get_type(self): return 'equity_trade' - def send_all(self): - trade_start = datetime.datetime.now() - minute = datetime.timedelta(minutes=1) - price = random.uniform(5.0, 50.0) - - for i in range(self.count): - if not self.sync.confirm(): - break - price = price + random.uniform(-0.05, 0.05) - event = {'sid':self.sid, - 'dt':qutil.format_date(trade_start + (minute * i)), - 'price':price, - 'volume':random.randrange(100,10000,100)} - self.send(event) + def do_work(self): + if(self.incr == self.count): + self.signal_done() + return + self.price = self.price + random.uniform(-0.05, 0.05) + event = {'sid':self.sid, + 'dt':qutil.format_date(self.trade_start + (self.minute * self.incr)), + 'price':self.price, + 'volume':random.randrange(100,10000,100)} + self.send(event) + self.incr += 1 diff --git a/zipline/test/client.py b/zipline/test/client.py index c04d3bd3..f13dfc4c 100644 --- a/zipline/test/client.py +++ b/zipline/test/client.py @@ -1,64 +1,44 @@ -from gevent_zeromq import zmq import json import zipline.util as qutil +import zipline.messaging as qmsg -class TestClient(object): +class TestClient(qmsg.Component): - def __init__(self, utest, expected_msg_count=0): - self.address = None - self.sync = None + def __init__(self, addresses, utest, expected_msg_count=0): + qmsg.Component.__init__(self, addresses) self.received_count = 0 self.expected_msg_count = expected_msg_count - self.error = False self.utest = utest - self.data_feed = None - self.context = None + self.prev_dt = None + + def get_id(self): + return "TEST_CLIENT" + + def open(self): + self.data_feed, self.poller = self.connect_result() - def run(self): - - try: - qutil.LOGGER.info("running the client") - self.context = zmq.Context() - - self.data_feed = self.context.socket(zmq.PULL) + def do_work(self): + socks = dict(self.poller.poll(2000)) #timeout after 2 seconds. + if self.data_feed in socks and socks[self.data_feed] == self.zmq.POLLIN: + msg = self.data_feed.recv() + if(self.is_done_message(msg)): + qutil.LOGGER.info("Client is DONE!") + self.signal_done() + self.utest.assertEqual(self.expected_msg_count, self.received_count, + "The client should have received ({n}) the same number of messages as the feed sent ({m})." + .format(n=self.received_count, m=self.expected_msg_count)) + return + + self.received_count += 1 + event = json.loads(msg) + if(self.prev_dt != None): + if(not event['dt'] >= self.prev_dt): + raise Exception("Message out of order: {date} after {prev}".format(date=event['dt'], prev=prev_dt)) - qutil.LOGGER.info("connecting to {address}".format(address=self.address)) - self.data_feed.connect(self.address) - - self.sync.open() - - self.poller = zmq.Poller() - self.poller.register(self.data_feed, zmq.POLLIN) + self.prev_dt = event['dt'] + if(self.received_count % 100 == 0): + qutil.LOGGER.info("received {n} messages".format(n=self.received_count)) - qutil.LOGGER.info("Starting the client loop") - - prev_dt = None - while self.sync.confirm(): - socks = dict(self.poller.poll(2000)) #timeout after 2 seconds. - if self.data_feed in socks and socks[self.data_feed] == zmq.POLLIN: - msg = self.data_feed.recv() - if(msg == "DONE"): - qutil.LOGGER.info("DONE!") - break - self.received_count += 1 - event = json.loads(msg) - if(prev_dt != None): - if(not event['dt'] >= prev_dt): - raise Exception("Message out of order: {date} after {prev}".format(date=event['dt'], prev=prev_dt)) - - prev_dt = event['dt'] - if(self.received_count % 100 == 0): - qutil.LOGGER.info("received {n} messages".format(n=self.received_count)) - - qutil.LOGGER.info("received {n} messages".format(n=self.received_count)) - self.data_feed.close() - self.sync.close() - except: - self.error = True - qutil.LOGGER.exception("Error in test client.") - finally: - self.context.destroy() + - self.utest.assertEqual(self.expected_msg_count, self.received_count, - "The client should have received ({n}) the same number of messages as the feed sent ({m})." - .format(n=self.received_count, m=self.expected_msg_count)) \ No newline at end of file + \ No newline at end of file diff --git a/zipline/test/test_messaging.py b/zipline/test/test_messaging.py index f0450655..9e1d5747 100644 --- a/zipline/test/test_messaging.py +++ b/zipline/test/test_messaging.py @@ -7,7 +7,7 @@ import unittest2 as unittest import multiprocessing import time -from zipline.core import ThreadSimulator, ProcessSimulator, DataFeed +from zipline.simulator import ThreadSimulator, ProcessSimulator from zipline.transforms.technical import MovingAverage from zipline.sources import RandomEquityTrades import zipline.util as qutil @@ -17,7 +17,7 @@ from zipline.test.client import TestClient qutil.configure_logging() -class MessagingTestCase(unittest.TestCase): +class ThreadSimulatorTestCase(unittest.TestCase): """Tests the message passing: datasources -> feed -> transforms -> merge -> client""" def setUp(self): @@ -29,38 +29,40 @@ class MessagingTestCase(unittest.TestCase): 'result_address' : "tcp://127.0.0.1:{port}".format(port=10104) } + def get_simulator(self): - return ProcessSimulator() + return ThreadSimulator(self.addresses) - def dtest_sources_only(self): + + def test_sources_only(self): """streams events from two data sources, no transforms.""" sim = self.get_simulator() - ret1 = RandomEquityTrades(133, "ret1", 400) - ret2 = RandomEquityTrades(134, "ret2", 400) - client = TestClient(self, expected_msg_count=800) + ret1 = RandomEquityTrades(133, "ret1", 400, self.addresses) + ret2 = RandomEquityTrades(134, "ret2", 400, self.addresses) + client = TestClient(self.addresses, self, expected_msg_count=800) sim.register_components([ret1, ret2, client]) sim.simulate() - self.assertEqual(sim.feed.data_buffer.pending_messages(), 0, + self.assertEqual(sim.feed.pending_messages(), 0, "The feed should be drained of all messages, found {n} remaining." - .format(n=sim.feed.data_buffer.pending_messages())) + .format(n=sim.feed.pending_messages())) - def test_merged_to_client(self): + def test_transforms(self): """ 2 datasources -> feed -> 2 moving average transforms -> transform merge -> testclient verify message count at client. """ sim = self.get_simulator() - ret1 = RandomEquityTrades(133, "ret1", 5000) - ret2 = RandomEquityTrades(134, "ret2", 5000) + ret1 = RandomEquityTrades(133, "ret1", 5000, self.addresses) + ret2 = RandomEquityTrades(134, "ret2", 5000, self.addresses) mavg1 = MovingAverage("mavg1", 30) mavg2 = MovingAverage("mavg2", 60) - client = TestClient(self, expected_msg_count=10000) - sim.register_components[ret1, ret2, mavg1, mavg2, client] + client = TestClient(self.addresses, self, expected_msg_count=10000) + sim.register_components([ret1, ret2, mavg1, mavg2, client]) sim.simulate() - self.assertEqual(sim.feed.data_buffer.pending_messages(), 0, "The feed should be drained of all messages.") + self.assertEqual(sim.feed.pending_messages(), 0, "The feed should be drained of all messages.") def dtest_error_in_feed(self): ret1 = RandomEquityTrades(133, "ret1", 400) @@ -69,19 +71,14 @@ class MessagingTestCase(unittest.TestCase): mavg1 = MovingAverage("mavg1", 30) mavg2 = MovingAverage("mavg2", 60) transforms = {"mavg1":mavg1, "mavg2":mavg2} - client = TestClient(self, expected_msg_count=0) + client = TestClient(self.addresses, self, expected_msg_count=0) sim = self.get_simulator(sources, transforms, client) sim.feed = DataFeedErr(sources.keys(), sim.data_address, sim.feed_address, sim.performance_address, qmsg.Sync(sim, "DataFeedErrorGenerator")) sim.simulate() -class DataFeedErr(DataFeed): - """Helper class for testing, simulates exceptions inside the DataFeed""" - def __init__(self, source_list, data_address, feed_address, perf_address, sync): - DataFeed.__init__(self, source_list, data_address, feed_address, perf_address, sync) +class ProcessSimulatorTestCase(ThreadSimulatorTestCase): - def handle_all(self): - #time.sleep(1000) - raise Exception("simulated error in data feed from test helper") - - + def get_simulator(self): + return ProcessSimulator(self.addresses) + diff --git a/zipline/transforms/technical.py b/zipline/transforms/technical.py index fc19bd76..dcefc87e 100644 --- a/zipline/transforms/technical.py +++ b/zipline/transforms/technical.py @@ -5,7 +5,7 @@ TODO: add trailing stop """ import datetime -from zipline.core import BaseTransform +from zipline.messaging import BaseTransform import zipline.util as qutil class MovingAverage(BaseTransform):