Rework imports on tests.

This commit is contained in:
Stephen Diehl
2012-05-14 11:17:56 -04:00
parent fc7cf23397
commit 73557b907f
3 changed files with 85 additions and 85 deletions
+31 -31
View File
@@ -18,7 +18,7 @@ from zipline.simulator import AddressAllocator
from zipline.lines import SimulatedTrading
from zipline.finance.performance import PerformanceTracker
from zipline.utils.protocol_utils import namedict
from zipline.finance.trading import SIMULATION_STYLE
from zipline.finance.trading import TransactionSimulator, SIMULATION_STYLE
DEFAULT_TIMEOUT = 15 # seconds
EXTENDED_TIMEOUT = 90
@@ -411,7 +411,7 @@ class FinanceTestCase(TestCase):
alternate = params.get('alternate')
# if present, expect transaction amounts to match orders exactly.
complete_fill = params.get('complete_fill')
trading_environment = factory.create_trading_environment()
trade_sim = TransactionSimulator()
price = [10.1] * trade_count
@@ -419,19 +419,19 @@ class FinanceTestCase(TestCase):
start_date = trading_environment.first_open
sid = 1
generated_trades = factory.create_trade_history(
sid,
price,
volume,
trade_interval,
trading_environment
generated_trades = factory.create_trade_history(
sid,
price,
volume,
trade_interval,
trading_environment
)
if alternate:
alternator = -1
else:
alternator = 1
order_date = start_date
for i in xrange(order_count):
order = namedict(
@@ -443,7 +443,7 @@ class FinanceTestCase(TestCase):
})
trade_sim.add_open_order(order)
order_date = order_date + order_interval
# move after market orders to just after market next
# market open.
@@ -451,40 +451,40 @@ class FinanceTestCase(TestCase):
if order_date.minute >= 00:
order_date = order_date + timedelta(days=1)
order_date = order_date.replace(hour=14, minute=30)
# there should now be one open order list stored under the sid
oo = trade_sim.open_orders
self.assertEqual(len(oo), 1)
self.assertTrue(oo.has_key(sid))
order_list = oo[sid]
self.assertEqual(order_count, len(order_list))
for i in xrange(order_count):
order = order_list[i]
self.assertEqual(order.sid, sid)
self.assertEqual(order.amount, order_amount * alternator**i)
tracker = PerformanceTracker(trading_environment)
# this approximates the loop inside TradingSimulationClient
transactions = []
for trade in generated_trades:
if trade_delay:
trade.dt = trade.dt + trade_delay
txn = trade_sim.apply_trade_to_open_orders(trade)
if txn:
transactions.append(txn)
trade.TRANSACTION = txn
transactions.append(txn)
trade.TRANSACTION = txn
else:
trade.TRANSACTION = None
tracker.process_event(trade)
tracker.process_event(trade)
if complete_fill:
self.assertEqual(len(transactions), len(order_list))
self.assertEqual(len(transactions), len(order_list))
total_volume = 0
for i in xrange(len(transactions)):
txn = transactions[i]
@@ -492,18 +492,18 @@ class FinanceTestCase(TestCase):
if complete_fill:
order = order_list[i]
self.assertEqual(order.amount, txn.amount)
self.assertEqual(total_volume, expected_txn_volume)
self.assertEqual(total_volume, expected_txn_volume)
self.assertEqual(len(transactions), expected_txn_count)
cumulative_pos = tracker.cumulative_performance.positions[sid]
self.assertEqual(total_volume, cumulative_pos.amount)
# the open orders should now be empty
oo = trade_sim.open_orders
self.assertTrue(oo.has_key(sid))
order_list = oo[sid]
self.assertEqual(0, len(order_list))
+1 -1
View File
@@ -13,7 +13,7 @@ import zipline.utils.factory as factory
from zipline.utils import logger
import zipline.protocol as zp
from zipline.sources import SpecificEquityTrades
from zipline.finance.sources import SpecificEquityTrades
DEFAULT_TIMEOUT = 5 # seconds
+53 -53
View File
@@ -24,10 +24,10 @@ def load_market_data():
# second=0,
# tzinfo=pytz.utc
#)
daily_return = risk.DailyReturn(date=event_dt, returns=returns)
bm_returns.append(daily_return)
bm_returns = sorted(bm_returns, key=lambda(x): x.date)
bm_returns = sorted(bm_returns, key=lambda(x): x.date)
fp_tr = open(".//tests/treasury_curves.msgpack", "rb")
tr_list = msgpack.loads(fp_tr.read())
tr_curves = {}
@@ -35,9 +35,9 @@ def load_market_data():
tr_dt = zp.tuple_to_date(packed_date)
#tr_dt = tr_dt.replace(hour=0, minute=0, second=0, tzinfo=pytz.utc)
tr_curves[tr_dt] = curve
return bm_returns, tr_curves
def create_trading_environment(year=2006):
"""Construct a complete environment with reasonable defaults"""
benchmark_returns, treasury_curves = load_market_data()
@@ -51,8 +51,9 @@ def create_trading_environment(year=2006):
period_end = end,
capital_base = 100000.0
)
return trading_environment
def create_trade(sid, price, amount, datetime):
row = zp.namedict({
'source_id' : "test_factory",
@@ -70,7 +71,7 @@ def get_next_trading_dt(current, interval, trading_calendar):
next = next + interval
if trading_calendar.is_market_hours(next):
break
return next
def create_trade_history(sid, prices, amounts, interval, trading_calendar):
@@ -78,7 +79,7 @@ def create_trade_history(sid, prices, amounts, interval, trading_calendar):
current = trading_calendar.first_open
for price, amount in zip(prices, amounts):
trade = create_trade(sid, price, amount, current)
trades.append(trade)
current = get_next_trading_dt(current, interval, trading_calendar)
@@ -88,10 +89,10 @@ def create_trade_history(sid, prices, amounts, interval, trading_calendar):
def create_txn(sid, price, amount, datetime, btrid=None):
txn = zp.namedict({
'sid':sid,
'amount':amount,
'dt':datetime,
'price':price,
'sid' : sid,
'amount' : amount,
'dt' : datetime,
'price' : price,
})
return txn
@@ -115,15 +116,15 @@ def create_returns(daycount, trading_calendar):
test_range = []
current = trading_calendar.first_open
one_day = timedelta(days = 1)
for day in range(daycount):
for day in range(daycount):
current = current + one_day
if trading_calendar.is_trading_day(current):
r = risk.DailyReturn(current, random.random())
test_range.append(r)
return test_range
def create_returns_from_range(trading_calendar):
current = trading_calendar.first_open
@@ -134,53 +135,53 @@ def create_returns_from_range(trading_calendar):
r = risk.DailyReturn(current, random.random())
test_range.append(r)
current = get_next_trading_dt(current, one_day, trading_calendar)
return test_range
def create_returns_from_list(returns, trading_calendar):
current = trading_calendar.first_open
one_day = timedelta(days = 1)
test_range = []
#sometimes the range starts with a non-trading day.
if not trading_calendar.is_trading_day(current):
current = get_next_trading_dt(current, one_day, trading_calendar)
for return_val in returns:
for return_val in returns:
r = risk.DailyReturn(current, return_val)
test_range.append(r)
current = get_next_trading_dt(current, one_day, trading_calendar)
return test_range
def create_random_trade_source(sid, trade_count, trading_environment):
# create the source
source = RandomEquityTrades(sid, "rand-"+str(sid), trade_count)
# make the period_end of trading_environment match
cur = trading_environment.first_open
one_day = timedelta(days = 1)
for i in range(trade_count + 2):
cur = get_next_trading_dt(cur, one_day, trading_environment)
trading_environment.period_end = cur
return source
def create_daily_trade_source(sids, trade_count, trading_environment):
"""
creates trade_count trades for each sid in sids list.
first trade will be on trading_environment.period_start, and daily
thereafter for each sid. Thus, two sids should result in two trades per
day.
creates trade_count trades for each sid in sids list.
first trade will be on trading_environment.period_start, and daily
thereafter for each sid. Thus, two sids should result in two trades per
day.
Important side-effect: trading_environment.period_end will be modified
to match the day of the final trade.
to match the day of the final trade.
"""
return create_trade_source(
sids,
trade_count,
timedelta(days=1),
sids,
trade_count,
timedelta(days=1),
trading_environment
)
@@ -188,18 +189,18 @@ def create_daily_trade_source(sids, trade_count, trading_environment):
def create_minutely_trade_source(sids, trade_count, trading_environment):
"""
creates trade_count trades for each sid in sids list.
first trade will be on trading_environment.period_start, and every minute
thereafter for each sid. Thus, two sids should result in two trades per
minute.
creates trade_count trades for each sid in sids list.
first trade will be on trading_environment.period_start, and every minute
thereafter for each sid. Thus, two sids should result in two trades per
minute.
Important side-effect: trading_environment.period_end will be modified
to match the day of the final trade.
to match the day of the final trade.
"""
return create_trade_source(
sids,
trade_count,
timedelta(minutes=1),
sids,
trade_count,
timedelta(minutes=1),
trading_environment
)
@@ -210,22 +211,21 @@ def create_trade_source(sids, trade_count, trade_time_increment, trading_environ
volume = [100] * trade_count
start_date = trading_environment.first_open
generated_trades = create_trade_history(
sid,
price,
volume,
trade_time_increment,
trading_environment
generated_trades = create_trade_history(
sid,
price,
volume,
trade_time_increment,
trading_environment
)
trade_history.extend(generated_trades)
trade_history = sorted(trade_history, key=lambda(x): x.dt)
#set the trading environment's end to same dt as the last trade in the
#history.
trading_environment.period_end = trade_history[-1].dt
source = SpecificEquityTrades("flat", trade_history)
return source