mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-15 11:22:18 +08:00
71 lines
2.1 KiB
Python
71 lines
2.1 KiB
Python
import zmq
|
|
import zipline.protocol as zp
|
|
|
|
def drain_zipline(test, zipline):
|
|
assert test.ctx, "method expects a valid zmq context"
|
|
assert test.zipline_test_config, "method expects a valid test config"
|
|
assert isinstance(test.zipline_test_config, dict)
|
|
assert test.zipline_test_config['results_socket'], \
|
|
"need to specify a socket address for logs/perf/risk"
|
|
test.receiver = test.ctx.socket(zmq.PULL)
|
|
test.receiver.bind(test.zipline_test_config['results_socket'])
|
|
|
|
# start the simulation
|
|
zipline.simulate(blocking=False)
|
|
|
|
output = []
|
|
transaction_count = 0
|
|
while True:
|
|
msg = test.receiver.recv()
|
|
if msg == str(zp.CONTROL_PROTOCOL.DONE):
|
|
break
|
|
elif msg == "EXCEPTION":
|
|
output.append(msg)
|
|
break
|
|
else:
|
|
update = zp.BT_UPDATE_UNFRAME(msg)
|
|
output.append(update)
|
|
if update['prefix'] == 'PERF':
|
|
transaction_count += \
|
|
len(update['payload']['daily_perf']['transactions'])
|
|
|
|
del test.receiver
|
|
|
|
# some processes will exit after the message stream is
|
|
# finished. We block here to avoid collisions with subsequent
|
|
# ziplines.
|
|
for process in zipline.sim.subprocesses:
|
|
process.join()
|
|
|
|
return output, transaction_count
|
|
|
|
|
|
def assert_single_position(test, zipline):
|
|
output, transaction_count = drain_zipline(test, zipline)
|
|
|
|
test.assertTrue(zipline.sim.ready())
|
|
test.assertFalse(zipline.sim.exception)
|
|
|
|
test.assertEqual(
|
|
test.zipline_test_config['order_count'],
|
|
transaction_count
|
|
)
|
|
|
|
# the final message is the risk report, the second to
|
|
# last is the final day's results. Positions is a list of
|
|
# dicts.
|
|
closing_positions = output[-2]['payload']['daily_perf']['positions']
|
|
|
|
test.assertEqual(
|
|
len(closing_positions),
|
|
1,
|
|
"Portfolio should have one position."
|
|
)
|
|
|
|
sid = test.zipline_test_config['sid']
|
|
test.assertEqual(
|
|
closing_positions[0]['sid'],
|
|
sid,
|
|
"Portfolio should have one position in " + str(sid)
|
|
)
|