From 73557b907f3bc2751669e4e13a7ea36277479712 Mon Sep 17 00:00:00 2001 From: Stephen Diehl Date: Mon, 14 May 2012 11:17:56 -0400 Subject: [PATCH] Rework imports on tests. --- tests/test_finance.py | 62 +++++++++++------------ tests/test_protocol.py | 2 +- zipline/utils/factory.py | 106 +++++++++++++++++++-------------------- 3 files changed, 85 insertions(+), 85 deletions(-) diff --git a/tests/test_finance.py b/tests/test_finance.py index ba9eb9af..0aa14bba 100644 --- a/tests/test_finance.py +++ b/tests/test_finance.py @@ -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)) - - - + + + diff --git a/tests/test_protocol.py b/tests/test_protocol.py index ec39a352..2250c4c0 100644 --- a/tests/test_protocol.py +++ b/tests/test_protocol.py @@ -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 diff --git a/zipline/utils/factory.py b/zipline/utils/factory.py index a19b9a48..250ca670 100644 --- a/zipline/utils/factory.py +++ b/zipline/utils/factory.py @@ -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 -