Merge pull request #17 from quantopian/api

Api

zipline/component.py

@@ -448,6 +448,13 @@ class Component(object):
 
         """
         raise NotImplementedError
+        
+    @property
+    def is_blocking(self):
+        """
+        True if a zipline be held open for this component.
+        """
+        return False
 
     @property
     def get_pure(self):

zipline/finance/performance.py

@@ -185,7 +185,7 @@ class PerformanceTracker():
     def to_dict(self):
         """
         Creates a dictionary representing the state of this tracker.
-        Returns a dict object of the form:
+        Returns a dict object of the form described in header comments.
         """
 
         returns_list = [x.to_dict() for x in self.returns]
@@ -295,8 +295,8 @@ class PerformanceTracker():
         #if self.result_stream:
             ## TODO: proper framing
             #self.result_stream.send_pyobj(self.risk_report.to_dict())
-
-        self.result_stream.send_pyobj(None)
+        if self.result_stream:
+            self.result_stream.send_pyobj(None)
 
     def round_to_nearest(self, x, base=5):
         return int(base * round(float(x)/base))
@@ -368,6 +368,8 @@ class PerformancePeriod():
         #cash balance at start of period
         self.starting_cash          = starting_cash
         self.ending_cash            = starting_cash
+        
+        self.calculate_performance()
 
     def calculate_performance(self):
         self.ending_value = self.calculate_positions_value()

zipline/finance/risk.py

@@ -374,48 +374,3 @@ class RiskReport():
         if len(col) == 1:
             return col[0]
         return None
-
-
-class TradingEnvironment(object):
-
-    def __init__(
-        self, 
-        benchmark_returns, 
-        treasury_curves, 
-        period_start=None, 
-        period_end=None, 
-        capital_base=None,
-        frame_index=None
-    ):
-    
-        self.trading_days = []
-        self.trading_day_map = {}
-        self.treasury_curves = treasury_curves
-        self.benchmark_returns = benchmark_returns
-        self.frame_index = frame_index
-        self.period_start = period_start
-        self.period_end = period_end
-        self.capital_base = capital_base
-        for bm in benchmark_returns:
-            self.trading_days.append(bm.date)
-            self.trading_day_map[bm.date] = bm
-
-    def normalize_date(self, test_date):
-        return datetime.datetime(
-            year=test_date.year,
-            month=test_date.month,
-            day=test_date.day,
-            tzinfo=pytz.utc
-        )
-
-    def is_trading_day(self, test_date):
-        dt = self.normalize_date(test_date)
-        return self.trading_day_map.has_key(dt)
-
-    def get_benchmark_daily_return(self, test_date):
-        date = self.normalize_date(test_date)
-        if self.trading_day_map.has_key(date):
-            return self.trading_day_map[date].returns
-        else:
-            return 0.0
-

zipline/finance/trading.py

@@ -29,6 +29,12 @@ class TradeSimulationClient(qmsg.Component):
         )
         
         self.perf = perf.PerformanceTracker(self.trading_environment)
+        ##################################################################
+        # TODO: the next line of code need refactoring from RealDiehl
+        # The below sets up the performance object to trigger a full risk
+        # report with rolling periods over the entire test duration. We
+        # would prefer something more explicit than a callback.
+        ##################################################################
         self.on_done = self.perf.handle_simulation_end
         
     
@@ -61,6 +67,7 @@ class TradeSimulationClient(qmsg.Component):
             if msg == str(zp.CONTROL_PROTOCOL.DONE):
                 qutil.LOGGER.info("Client is DONE!")
                 self.run_callbacks()
+                self.signal_order_done()
                 self.signal_done()
                 return
             
@@ -107,6 +114,13 @@ class TradeSimulationClient(qmsg.Component):
         self.order_socket.send(str(zp.ORDER_PROTOCOL.DONE))
         
     def queue_event(self, event):
+        ##################################################################
+        # TODO: the next line of code need refactoring from RealDiehl
+        # the performance class needs to process each event, without skipping
+        # and any callbacks should wait until the performance has been 
+        # updated, so that down stream components can safely assume that
+        # performance is up to date.
+        ##################################################################
         self.perf.process_event(event)
         if self.event_queue == None:
             self.event_queue = []
@@ -140,6 +154,14 @@ class OrderDataSource(qmsg.DataSource):
     def get_type(self):
         return zp.DATASOURCE_TYPE.ORDER
         
+    #
+    @property
+    def is_blocking(self):
+        """
+        This datasource is in a loop with the TradingSimulationClient
+        """
+        return False
+        
     def open(self):
         qmsg.DataSource.open(self)
         self.order_socket = self.bind_order()
@@ -157,14 +179,14 @@ class OrderDataSource(qmsg.DataSource):
         orders = []
         count = 0
         while True:
-            
+                        
             (rlist, wlist, xlist) = select(
                 [self.order_socket],
                 [],
                 [self.order_socket],
                 #allow half the time of a heartbeat for the order
                 #timeout, so we have time to signal we are done.
-                timeout=self.heartbeat_timeout/2000
+                #timeout=self.heartbeat_timeout/2000
             ) 
         
             
@@ -173,6 +195,7 @@ class OrderDataSource(qmsg.DataSource):
                 #no order message means there was a timeout above, 
                 #and the client is done sending orders (but isn't
                 #telling us himself!).
+                qutil.LOGGER.warn("signaling orders done on timeout.")
                 self.signal_done()
                 return
                 
@@ -303,6 +326,57 @@ class TransactionSimulator(qmsg.BaseTransform):
                 }
         return zp.namedict(txn) 
                 
-                
+
+class TradingEnvironment(object):
+
+    def __init__(
+        self, 
+        benchmark_returns, 
+        treasury_curves, 
+        period_start=None, 
+        period_end=None, 
+        capital_base=None
+    ):
+    
+        self.trading_days = []
+        self.trading_day_map = {}
+        self.treasury_curves = treasury_curves
+        self.benchmark_returns = benchmark_returns
+        self.frame_index = ['sid', 'volume', 'dt', 'price', 'changed']
+        self.period_start = period_start
+        self.period_end = period_end
+        self.capital_base = capital_base
+            
+        for bm in benchmark_returns:
+            self.trading_days.append(bm.date)
+            self.trading_day_map[bm.date] = bm
+
+    def normalize_date(self, test_date):
+        return datetime.datetime(
+            year=test_date.year,
+            month=test_date.month,
+            day=test_date.day,
+            tzinfo=pytz.utc
+        )
+
+    def is_trading_day(self, test_date):
+        dt = self.normalize_date(test_date)
+        return self.trading_day_map.has_key(dt)
+
+    def get_benchmark_daily_return(self, test_date):
+        date = self.normalize_date(test_date)
+        if self.trading_day_map.has_key(date):
+            return self.trading_day_map[date].returns
+        else:
+            return 0.0
+            
+    def add_to_frame(self, name):
+        """
+        Add an entry to the frame index. 
+        :param name: new index entry name. Used by TradingSimulationClient
+        to 
+        """
+        self.frame_index.append(name)
+
                 
 

zipline/lines.py

@@ -0,0 +1,338 @@
+"""
+Ziplines are composed of multiple components connected by asynchronous 
+messaging. All ziplines follow a general topology of parallel sources, 
+datetimestamp serialization, parallel transformations, and finally sinks. 
+Furthermore, many ziplines have common needs. For example, all trade 
+simulations require a 
+:py:class:`~zipline.finance.trading.TradeSimulationClient`, an
+:py:class:`~zipline.finance.trading.OrderSource`, and a 
+:py:class:`~zipline.finance.trading.TransactionSimulator` (a transform).
+
+To establish best practices and minimize code replication, the lines module 
+provides complete zipline topologies. You can extend any zipline without
+the need to extend the class. Simply instantiate any additional components
+that you would like included in the zipline, and add them to the zipline 
+before invoking simulate. 
+
+        
+        Here is a diagram of the SimulatedTrading zipline:
+        
+        
+            +----------------------+  +------------------------+
+        +-->|  Orders DataSource   |  |    (DataSource added   |
+        |   |  Integrates algo     |  |     via add_source)    |
+        |   |  orders into history |  |                        |
+        |   +--------------------+-+  +-+----------------------+
+        |                        |      |
+        |                        |      |
+        |                        v      v
+        |                       +---------+
+        |                       |   Feed  |
+        |                       +-+------++
+        |                         |      |
+        |                         |      |    
+        |                         v      v
+        |    +----------------------+   +----------------------+
+        |    | Transaction          |   |                      |
+        |    | Transform simulates  |   |  (Transforms added   |
+        |    | trades based on      |   |   via add_transform) |
+        |    | orders from algo.    |   |                      |
+        |    +-------------------+--+   +-+--------------------+
+        |                        |        |
+        |                        |        |
+        |                        v        v
+        |                      +------------+
+        |                      |    Merge   |
+        |                      +------+-----+
+        |                             |
+        |                             |
+        |                             V
+        |               +--------------------------------+
+        |               |                                |
+        |               |     TradingSimulationClient    |
+        |  orders       |     tracks performance and     |
+        +---------------+     provides API to algorithm. |
+                        |                                |
+                        +---------------------+----------+
+                                  ^           |
+                                  | orders    |  frames
+                                  |           |
+                                  |           v
+                        +---------+-----------------------+
+                        |                                 |
+                        |  Algorithm added via            |
+                        |  __init__.                      |
+                        |                                 |
+                        |                                 |
+                        |                                 |
+                        +---------------------------------+
+
+"""
+
+import mock
+import pytz
+
+from datetime import datetime, timedelta
+from collections import defaultdict
+
+from nose.tools import timed
+
+import zipline.test.factory as factory
+import zipline.util as qutil
+import zipline.finance.risk as risk
+import zipline.protocol as zp
+import zipline.finance.performance as perf
+import zipline.messaging as zmsg
+
+from zipline.test.client import TestAlgorithm
+from zipline.sources import SpecificEquityTrades
+from zipline.finance.trading import TransactionSimulator, OrderDataSource, \
+TradeSimulationClient
+from zipline.simulator import AddressAllocator, Simulator
+from zipline.monitor import Controller
+
+
+
+class SimulatedTrading(object):
+    """
+        Zipline with::
+            - _no_ data sources.
+            - Trade simulation client, which is available to send callbacks on
+            events and also accept orders to be simulated.
+            - An order data source, which will receive orders from the trade
+            simulation client, and feed them into the event stream to be 
+            serialized and order alongside all other data source events.
+            - transaction simulation transformation, which receives the order
+            events and estimates a theoretical execution price and volume.
+            
+        All components in this zipline are subject to heartbeat checks and
+        a control monitor, which can kill the entire zipline in the event of
+        exceptions in one of the components or an external request to end the
+        simulation.
+    """
+    
+    def __init__(self, **config):
+        """
+        :param config: a dict with the following required properties::
+            - algorithm: a class that follows the algorithm protocol. Must
+            have a handle_frame method that accepts a pandas.Dataframe of the 
+            current state of the simulation universe. Must have an order 
+            property which can be set equal to the order method of 
+            trading_client. (TODO: where should this protocol be documented?)
+            - trading_environment: an instance of
+            :py:class:`zipline.trading.TradingEnvironment`
+            - allocator: an instance of 
+            :py:class:`zipline.simulator.AddressAllocator`
+            - simulator_class: a :py:class:`zipline.messaging.ComponentHost` 
+            subclass (not an instance)
+        """
+        assert isinstance(config, dict)
+        self.algorithm = config['algorithm']
+        self.allocator = config['allocator']
+        self.trading_environment = config['trading_environment']
+        
+        self.leased_sockets = []
+        self.sim_context = None
+        
+        sockets = self.allocate_sockets(8)
+        addresses = {
+            'sync_address'   : sockets[0],
+            'data_address'   : sockets[1],
+            'feed_address'   : sockets[2],
+            'merge_address'  : sockets[3],
+            'result_address' : sockets[4],
+            'order_address'  : sockets[5]
+        }
+
+        self.con = Controller(
+            sockets[6],
+            sockets[7],
+            logging = qutil.LOGGER
+        )
+
+        
+        self.sim = config['simulator_class'](addresses)
+            
+        self.clients = {}
+        self.trading_client = TradeSimulationClient(self.trading_environment)
+        self.clients[self.trading_client.get_id] = self.trading_client
+        
+        # setup all sources
+        self.sources = {}
+        self.order_source = OrderDataSource()
+        self.sources[self.order_source.get_id] = self.order_source
+        
+        #setup transforms
+        self.transaction_sim = TransactionSimulator()
+        self.transforms = {}
+        self.transforms[self.transaction_sim.get_id] = self.transaction_sim
+        
+        #register all components
+        self.sim.register_components([
+            self.trading_client, 
+            self.order_source, 
+            self.transaction_sim 
+            ])
+            
+        self.sim.register_controller( self.con )
+        self.sim.on_done = self.shutdown()
+        self.started = False
+        
+        ##################################################################
+        #TODO: the next two lines of code need refactoring from RealDiehl
+        ##################################################################
+        #wire up a callback inside the algorithm to receive frames from the
+        #trading client
+        self.trading_client.add_event_callback(self.algorithm.handle_frame)
+        #register the trading_client's order method with the algorithm
+        self.algorithm.set_order(self.trading_client.order)
+    
+    @staticmethod
+    def create_test_zipline(**config):
+        """
+        :param config: A configuration object that is a dict with::
+            - environment - a \
+            :py:class:`zipline.finance.trading.TradeSimulationClient`
+            - allocator - a :py:class:`zipline.simulator.AddressAllocator`
+            - sid - an integer, which will be used as the security ID. 
+            - order_count - the number of orders the test algo will place,
+            defaults to 100
+            - trade_count - the number of trades to simulate, defaults to 100
+            - simulator_class - optional parameter that provides an alternative 
+            subclass of ComponentHost to hold the whole zipline. Defaults to
+            :py:class:`zipline.simulator.Simulator`   
+            - algorithm - optional parameter providing an algorithm. defaults
+            to :py:class:`zipline.test.client.TestAlgorithm`
+        """
+        assert isinstance(config, dict)
+        
+        allocator = config['allocator']
+        sid = config['sid']
+        
+        #--------------------
+        # Trading Environment
+        #--------------------
+        if config.has_key('environment'):
+            trading_environment = config['environment']
+        else:
+            trading_environment = factory.create_trading_environment()
+        
+        if config.has_key('order_count'):
+            order_count = config['order_count']
+        else:
+            order_count = 100
+            
+        if config.has_key('trade_count'):
+            trade_count = config['trade_count']
+        else:
+            trade_count = 100
+            
+        if config.has_key('simulator_class'):
+            simulator_class = config['simulator_class']
+        else:
+            simulator_class = Simulator
+              
+        #-------------------
+        # Trade Source
+        #-------------------
+        sids = [sid]
+        #-------------------
+        trade_source = factory.create_daily_trade_source(
+            sids,
+            trade_count,
+            trading_environment
+        )
+        #-------------------
+        # Create the Algo
+        #-------------------
+        if config.has_key('algorithm'):
+            test_algo = config['algorithm']
+        else:
+            order_amount = 100
+            test_algo = TestAlgorithm(
+                sid,
+                order_amount,
+                order_count
+            )
+        #-------------------
+        # Simulation
+        #-------------------
+        zipline = SimulatedTrading(**{
+            'algorithm':test_algo,
+            'trading_environment':trading_environment,
+            'allocator':allocator,
+            'simulator_class':simulator_class
+        })
+        #-------------------
+
+        zipline.add_source(trade_source)
+
+        return zipline
+        
+    def add_source(self, source):
+        assert isinstance(source, zmsg.DataSource)
+        self.check_started()    
+        self.sim.register_components([source])
+        self.sources[source.get_id] = source
+        
+    
+    def add_transform(self, transform):
+        assert isinstance(transform, zmsg.BaseTransform)
+        self.check_started()
+        self.sim.register_components([transform])
+        self.sources[transform.get_id] = transform
+    
+    def check_started(self):
+        if self.started:
+            raise ZiplineException("TradeSimulation", "You cannot add sources \
+            after the simulation has begun.")
+    
+    def get_cumulative_performance(self):
+        return self.trading_client.perf.cumulative_performance.to_dict()
+    
+    def allocate_sockets(self, n):
+        """
+        Allocate sockets local to this line, track them so
+        we can gc after test run.
+        """
+
+        assert isinstance(n, int)
+        assert n > 0
+
+        leased = self.allocator.lease(n)
+
+        self.leased_sockets.extend(leased)
+        return leased
+        
+    def simulate(self, blocking=False):
+        self.started = True
+        self.sim_context = self.sim.simulate()
+        if blocking:
+            self.sim_context.join()
+
+    def shutdown(self):
+        self.allocator.reaquire(*self.leased_sockets)
+        
+    #--------------------------------#
+    # Component property accessors   #
+    #--------------------------------#
+    
+    def get_positions(self):
+        """
+        returns current positions as a dict. draws from the cumulative
+        performance period in the performance tracker.
+        """
+        perf = self.trading_client.perf.cumulative_performance
+        positions = perf.get_positions()
+        return positions
+        
+class ZiplineException(Exception):
+    def __init__(self, zipline_name, msg):
+        self.name = zipline_name
+        self.message = msg
+        
+    def __str__(self):
+        return "Unexpected exception {line}: {msg}".format(
+            line=self.name, 
+            msg=self.message
+        )

zipline/messaging.py

@@ -65,7 +65,6 @@ class ComponentHost(Component):
         communication with them.
         """
         assert isinstance(components, list)
-
         for component in components:
 
             component.gevent_needed = self.gevent_needed
@@ -78,9 +77,13 @@ class ComponentHost(Component):
 
             if isinstance(component, DataSource):
                 self.feed.add_source(component.get_id)
+                if not component.is_blocking:
+                    self.feed.ds_finished_counter +=1 
             if isinstance(component, BaseTransform):
                 self.merge.add_source(component.get_id)
-
+                if not component.is_blocking:
+                    self.feed.ds_finished_counter +=1
+                        
     def unregister_component(self, component_id):
         del self.components[component_id]
         del self.sync_register[component_id]
@@ -220,6 +223,7 @@ class ParallelBuffer(Component):
 
                 if len(self.data_buffer) == self.ds_finished_counter:
                     #drain any remaining messages in the buffer
+                    qutil.LOGGER.debug("draining feed")
                     self.drain()
                     self.signal_done()
             else:
@@ -261,7 +265,7 @@ class ParallelBuffer(Component):
         """
         Send the (chronologically) next message in the buffer.
         """
-        if(not(self.is_full() or self.draining)):
+        if not (self.is_full() or self.draining):
             return
 
         event = self.next()
@@ -433,6 +437,10 @@ class BaseTransform(Component):
     def get_type(self):
         return COMPONENT_TYPE.CONDUIT
 
+    @property
+    def is_blocking(self):
+        return True
+
     def open(self):
         """
         Establishes zmq connections.
@@ -550,6 +558,10 @@ class DataSource(Component):
     @property
     def get_id(self):
         return self.id
+        
+    @property
+    def is_blocking(self):
+        return True
 
     @property
     def get_type(self):

zipline/simulator.py

@@ -33,6 +33,10 @@ class Simulator(ComponentHost):
         ComponentHost.__init__(self, addresses)
         self.subthreads = []
         self.running = False
+    
+    @property
+    def get_id(self):
+        return 'Simple Simulator'
 
     def launch_controller(self):
         thread = threading.Thread(target=self.controller.run)
@@ -49,6 +53,7 @@ class Simulator(ComponentHost):
         self.running = True
 
         return thread
+        
 
     def did_clean_shutdown(self):
         return not any([t.isAlive() for t in self.subthreads])

zipline/test/client.py

@@ -87,15 +87,17 @@ class TestClient(qmsg.Component):
 
 class TestAlgorithm():
     
-    def __init__(self, sid, amount, order_count, trading_client):
-        self.trading_client = trading_client
-        self.trading_client.add_event_callback(self.handle_frame)
+    def __init__(self, sid, amount, order_count):
         self.count = order_count
         self.sid = sid
         self.amount = amount
         self.incr = 0
         self.done = False
-    
+        self.order = None
+        
+    def set_order(self, order_callable):
+        self.order = order_callable
+        
     def handle_frame(self, frame):
         for dt, s in frame.iteritems():     
             data = {}
@@ -103,8 +105,14 @@ class TestAlgorithm():
             event = zp.namedict(data)
             #place an order for 100 shares of sid:133
             if self.incr < self.count:
-                self.trading_client.order(self.sid, self.amount)
+                self.order(self.sid, self.amount)
                 self.incr += 1
-            elif not self.done:
-                self.trading_client.signal_order_done()
-                self.done = True
+                
+class NoopAlgorithm(object):
+        
+    def set_order(self, order_callable):
+        pass
+    
+    def handle_frame(self, frame):
+        pass
+        

zipline/test/factory.py

@@ -1,17 +1,23 @@
-import datetime
+"""
+Factory functions to prepare useful data for tests.
+"""
 import pytz
 import msgpack
 import random
+
+from datetime import datetime, timedelta
 import zipline.util as qutil
 import zipline.finance.risk as risk
 import zipline.protocol as zp
+from zipline.sources import SpecificEquityTrades
+from zipline.finance.trading import TradingEnvironment
 
 def load_market_data():
     fp_bm = open("./zipline/test/benchmark.msgpack", "rb")
     bm_map = msgpack.loads(fp_bm.read())
     bm_returns = []
     for epoch, returns in bm_map.iteritems():
-        event_dt = datetime.datetime.fromtimestamp(epoch)
+        event_dt = datetime.fromtimestamp(epoch)
         event_dt = event_dt.replace(
             hour=0, 
             minute=0, 
@@ -26,13 +32,26 @@ def load_market_data():
     tr_map = msgpack.loads(fp_tr.read())
     tr_curves = {}
     for epoch, curve in tr_map.iteritems():
-        tr_dt = datetime.datetime.fromtimestamp(epoch)
+        tr_dt = datetime.fromtimestamp(epoch)
         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():
+    """Construct a complete environment with reasonable defaults"""
+    benchmark_returns, treasury_curves = load_market_data()
 
+    start = datetime.strptime("01/01/2006","%m/%d/%Y")
+    start = start.replace(tzinfo=pytz.utc)
+    trading_environment = TradingEnvironment(
+        benchmark_returns,
+        treasury_curves,
+        period_start = start,
+        capital_base = 100000.0
+    )
+    
+    return trading_environment
 def create_trade(sid, price, amount, datetime):
     row = zp.namedict({
         'source_id' : "test_factory",
@@ -57,7 +76,7 @@ def create_trade_history(sid, prices, amounts, start_time, interval, trading_cal
 
             current = current + interval
         else:
-            current = current + datetime.timedelta(days=1)
+            current = current + timedelta(days=1)
 
     return trades
 
@@ -81,7 +100,7 @@ def create_txn_history(sid, priceList, amtList, startTime, interval, trading_cal
             current = current + interval
 
         else:
-            current = current + datetime.timedelta(days=1)
+            current = current + timedelta(days=1)
 
     return txns
 
@@ -90,7 +109,7 @@ def create_returns(daycount, start, trading_calendar):
     i = 0
     test_range = []
     current = start.replace(tzinfo=pytz.utc)
-    one_day = datetime.timedelta(days = 1)
+    one_day = timedelta(days = 1)
     while i < daycount: 
         i += 1
         r = risk.DailyReturn(current, random.random())
@@ -102,7 +121,7 @@ def create_returns(daycount, start, trading_calendar):
 def create_returns_from_range(start, end, trading_calendar):
     current = start.replace(tzinfo=pytz.utc)
     end = end.replace(tzinfo=pytz.utc)
-    one_day = datetime.timedelta(days = 1)
+    one_day = timedelta(days = 1)
     test_range = []
     i = 0
     while current <= end: 
@@ -117,7 +136,7 @@ def create_returns_from_range(start, end, trading_calendar):
     
 def create_returns_from_list(returns, start, trading_calendar):
     current = start.replace(tzinfo=pytz.utc)
-    one_day = datetime.timedelta(days = 1)
+    one_day = timedelta(days = 1)
     test_range = []
     i = 0
     while len(test_range) < len(returns): 
@@ -128,3 +147,40 @@ def create_returns_from_list(returns, start, trading_calendar):
         current = current + one_day
     return sorted(test_range, key=lambda(x):x.date)
 
+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. 
+    
+    Important side-effect: trading_environment.period_end will be modified
+    to match the day of the final trade. 
+    """
+    trade_history = []
+    for sid in sids:
+        price = [10.1] * trade_count
+        volume = [100] * trade_count
+        start_date = trading_environment.period_start
+        trade_time_increment = timedelta(days=1)
+
+        generated_trades = create_trade_history( 
+            sid, 
+            price, 
+            volume, 
+            start_date, 
+            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
+        

zipline/test/test_finance.py

@@ -17,9 +17,10 @@ import zipline.finance.performance as perf
 from zipline.test.client import TestAlgorithm
 from zipline.sources import SpecificEquityTrades
 from zipline.finance.trading import TransactionSimulator, OrderDataSource, \
-TradeSimulationClient
+TradeSimulationClient, TradingEnvironment
 from zipline.simulator import AddressAllocator, Simulator
 from zipline.monitor import Controller
+from zipline.lines import SimulatedTrading
 
 DEFAULT_TIMEOUT = 5 # seconds
 
@@ -31,333 +32,78 @@ class FinanceTestCase(TestCase):
 
     def setUp(self):
         qutil.configure_logging()
-        self.benchmark_returns, self.treasury_curves = \
-        factory.load_market_data()
-
-        self.trading_environment = risk.TradingEnvironment(
-            self.benchmark_returns,
-            self.treasury_curves
-        )
-
-        self.allocator = allocator
-
-    def allocate_sockets(self, n):
-        """
-        Allocate sockets local to this test case, track them so
-        we can gc after test run.
-        """
-
-        assert isinstance(n, int)
-        assert n > 0
-
-        leased = self.allocator.lease(n)
-
-        self.leased_sockets[self.id()].extend(leased)
-        return leased
-
-    @timed(DEFAULT_TIMEOUT)
-    def test_trade_feed_protocol(self):
-
-        # TODO: Perhaps something more self-documenting for variables names?
-        sid = 133
-        price = [10.0] * 4
-        volume = [100] * 4
-
-        start_date = datetime.strptime("02/15/2012","%m/%d/%Y")
-        one_day_td = timedelta(days=1)
-
-        trades = factory.create_trade_history(
-            sid, 
-            price, 
-            volume, 
-            start_date, 
-            one_day_td, 
-            self.trading_environment
-        )
-
-        for trade in trades:
-            #simulate data source sending frame
-            msg = zp.DATASOURCE_FRAME(zp.namedict(trade))
-            #feed unpacking frame
-            recovered_trade = zp.DATASOURCE_UNFRAME(msg)
-            #feed sending frame
-            feed_msg = zp.FEED_FRAME(recovered_trade)
-            #transform unframing
-            recovered_feed = zp.FEED_UNFRAME(feed_msg)
-            #do a transform
-            trans_msg = zp.TRANSFORM_FRAME('helloworld', 2345.6)
-            #simulate passthrough transform -- passthrough shouldn't even
-            # unpack the msg, just resend.
-
-            passthrough_msg = zp.TRANSFORM_FRAME(zp.TRANSFORM_TYPE.PASSTHROUGH,\
-                    feed_msg)
-
-            #merge unframes transform and passthrough
-            trans_recovered = zp.TRANSFORM_UNFRAME(trans_msg)
-            pt_recovered = zp.TRANSFORM_UNFRAME(passthrough_msg)
-            #simulated merge
-            pt_recovered.PASSTHROUGH.merge(trans_recovered)
-            #frame the merged event
-            merged_msg = zp.MERGE_FRAME(pt_recovered.PASSTHROUGH)
-            #unframe the merge and validate values
-            event = zp.MERGE_UNFRAME(merged_msg)
-
-            #check the transformed value, should only be in event, not trade.
-            self.assertTrue(event.helloworld == 2345.6)
-            event.delete('helloworld')
-
-            self.assertEqual(zp.namedict(trade), event)
-
-    @timed(DEFAULT_TIMEOUT)
-    def test_order_protocol(self):
-        #client places an order
-        now = datetime.utcnow().replace(tzinfo=pytz.utc)
-        order = zp.namedict({
-            'dt':now,
-            'sid':133,
-            'amount':100
-        })
-        order_msg = zp.ORDER_FRAME(order)
-
-        #order datasource receives
-        order = zp.ORDER_UNFRAME(order_msg)
-        self.assertEqual(order.sid, 133)
-        self.assertEqual(order.amount, 100)
-        self.assertEqual(order.dt, now)
-        
-        #order datasource datasource frames the order
-        order_event = zp.namedict({
-            "sid"        : order.sid,
-            "amount"     : order.amount,
-            "dt"         : order.dt,
-            "source_id"  : zp.FINANCE_COMPONENT.ORDER_SOURCE,
-            "type"       : zp.DATASOURCE_TYPE.ORDER
-        })
-
-
-        order_ds_msg = zp.DATASOURCE_FRAME(order_event)
-
-        #transaction transform unframes
-        recovered_order = zp.DATASOURCE_UNFRAME(order_ds_msg)
-
-        self.assertEqual(now, recovered_order.dt)
-
-        #create a transaction from the order
-        txn = zp.namedict({
-            'sid'        : recovered_order.sid,
-            'amount'     : recovered_order.amount,
-            'dt'         : recovered_order.dt,
-            'price'      : 10.0,
-            'commission' : 0.50
-        })
-
-        #frame that transaction
-        txn_msg = zp.TRANSFORM_FRAME(zp.TRANSFORM_TYPE.TRANSACTION, txn)
-
-        #unframe
-        recovered_tx = zp.TRANSFORM_UNFRAME(txn_msg).TRANSACTION
-        self.assertEqual(recovered_tx.sid, 133)
-        self.assertEqual(recovered_tx.amount, 100)
+        self.zipline_test_config = {
+            'allocator':allocator,
+            'sid':133
+        }
 
     @timed(DEFAULT_TIMEOUT)
     def test_orders(self):
-
-        # Just verify sending and receiving orders.
-        # --------------
-
-        # Allocate sockets for the simulator components
-        sockets = self.allocate_sockets(8)
-
-        addresses = {
-            'sync_address'   : sockets[0],
-            'data_address'   : sockets[1],
-            'feed_address'   : sockets[2],
-            'merge_address'  : sockets[3],
-            'result_address' : sockets[4],
-            'order_address'  : sockets[5]
-        }
-
-        con = Controller(
-            sockets[6],
-            sockets[7],
-            logging = qutil.LOGGER
-        )
-
-        sim = Simulator(addresses)
-
-        # Simulation Components
-        # ---------------------
-
-        # TODO: Perhaps something more self-documenting for variables names?
-        sid = 133
-        price = [10.1] * 16
-        volume = [100] * 16
-        start_date = datetime.strptime("02/1/2012","%m/%d/%Y")
-        start_date = start_date.replace(tzinfo=pytz.utc)
-        trade_time_increment = timedelta(days=1)
-
-        trade_history = factory.create_trade_history( 
-            sid, 
-            price, 
-            volume, 
-            start_date, 
-            trade_time_increment, 
-            self.trading_environment 
-        )
-
-        set1 = SpecificEquityTrades("flat-133", trade_history)
-        
-        self.trading_environment.period_start = trade_history[0].dt
-        self.trading_environment.period_end = trade_history[-1].dt
-        self.trading_environment.capital_base = 10000
-        self.trading_environment.frame_index = ['sid', 'volume', 'dt', \
-        'price', 'changed']
-        
-        trading_client = TradeSimulationClient(self.trading_environment)
-        #client will send 10 orders for 100 shares of 133
-        test_algo = TestAlgorithm(133, 100, 10, trading_client)
-
-        order_source = OrderDataSource()
-        transaction_sim = TransactionSimulator()
-
-        sim.register_components([
-            trading_client, 
-            order_source, 
-            transaction_sim, 
-            set1
-        ])
-        sim.register_controller( con )
-
         # Simulation
         # ----------
-        sim_context = sim.simulate()
-        sim_context.join()
+        zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config)
+        zipline.simulate(blocking=True)
 
-        self.assertTrue(sim.ready())
-        self.assertFalse(sim.exception)
+        self.assertTrue(zipline.sim.ready())
+        self.assertFalse(zipline.sim.exception)
 
         # TODO: Make more assertions about the final state of the components.
-        self.assertEqual(sim.feed.pending_messages(), 0, \
+        self.assertEqual(zipline.sim.feed.pending_messages(), 0, \
             "The feed should be drained of all messages, found {n} remaining." \
-            .format(n=sim.feed.pending_messages()))
+            .format(n=zipline.sim.feed.pending_messages()))
 
 
     @timed(DEFAULT_TIMEOUT)
     def test_performance(self): 
-
-        # verify order -> transaction -> portfolio position.
-        # --------------
-
-        # Allocate sockets for the simulator components
-        sockets = self.allocate_sockets(8)
-
-        addresses = {
-            'sync_address'   : sockets[0],
-            'data_address'   : sockets[1],
-            'feed_address'   : sockets[2],
-            'merge_address'  : sockets[3],
-            'result_address' : sockets[4],
-            'order_address'  : sockets[5]
-        }
-
-        con = Controller(
-            sockets[6],
-            sockets[7],
-            logging = qutil.LOGGER
-        )
-
-        sim = Simulator(addresses)
-
-        # Simulation Components
-        # ---------------------
-
-        # TODO: Perhaps something more self-documenting for variables names?
-        trade_count = 100
-        sid = 133
-        price = [10.1] * trade_count
-        volume = [100] * trade_count
-        start_date = datetime.strptime("02/1/2012","%m/%d/%Y")
-        start_date = start_date.replace(tzinfo=pytz.utc)
-        trade_time_increment = timedelta(days=1)
-
-        trade_history = factory.create_trade_history( 
-            sid, 
-            price, 
-            volume, 
-            start_date, 
-            trade_time_increment, 
-            self.trading_environment 
-        )
-        
-        
-        self.trading_environment.period_start = trade_history[0].dt
-        self.trading_environment.period_end = trade_history[-1].dt
-        self.trading_environment.capital_base = 10000
-        self.trading_environment.frame_index = ['sid', 'volume', 'dt', \
-        'price', 'changed']
-        
-        set1 = SpecificEquityTrades("flat-133", trade_history)
-
-        #client sill send 10 orders for 100 shares of 133
-        trading_client = TradeSimulationClient(self.trading_environment)
-        test_algo = TestAlgorithm(133, 100, 10, trading_client)
-
-        order_source = OrderDataSource()
-        transaction_sim = TransactionSimulator()
-        
-        sim.register_components([
-            trading_client, 
-            order_source, 
-            transaction_sim, 
-            set1, 
-            ])
-        sim.register_controller( con )
-
-        # Simulation
-        # ----------
-        sim_context = sim.simulate()
-        sim_context.join()
+        #provide enough trades to ensure all orders are filled.
+        self.zipline_test_config['order_count'] = 100
+        self.zipline_test_config['trade_count'] = 200
+        zipline = SimulatedTrading.create_test_zipline(**self.zipline_test_config)
+        zipline.simulate(blocking=True)
 
         self.assertEqual(
-            sim.feed.pending_messages(), 
+            zipline.sim.feed.pending_messages(), 
             0, 
             "The feed should be drained of all messages, found {n} remaining." \
-            .format(n=sim.feed.pending_messages())
+            .format(n=zipline.sim.feed.pending_messages())
         )
         
         self.assertEqual(
-            sim.merge.pending_messages(), 
+            zipline.sim.merge.pending_messages(), 
             0, 
             "The merge should be drained of all messages, found {n} remaining." \
-            .format(n=sim.merge.pending_messages())
+            .format(n=zipline.sim.merge.pending_messages())
         )
 
         self.assertEqual(
-            test_algo.count,
-            test_algo.incr,
+            zipline.algorithm.count,
+            zipline.algorithm.incr,
             "The test algorithm should send as many orders as specified.")
             
+        order_source = zipline.sources[zp.FINANCE_COMPONENT.ORDER_SOURCE]
         self.assertEqual(
             order_source.sent_count, 
-            test_algo.count, 
+            zipline.algorithm.count, 
             "The order source should have sent as many orders as the algo."
         )
             
+        transaction_sim = zipline.transforms[zp.TRANSFORM_TYPE.TRANSACTION]
         self.assertEqual(
             transaction_sim.txn_count,
-            trading_client.perf.txn_count,
+            zipline.trading_client.perf.txn_count,
             "The perf tracker should handle the same number of transactions \
             as the simulator emits."
         ) 
         
         self.assertEqual(
-            len(trading_client.perf.cumulative_performance.positions), 
+            len(zipline.get_positions()), 
             1, 
             "Portfolio should have one position."
         )
         
+        SID = self.zipline_test_config['sid']
         self.assertEqual(
-            trading_client.perf.cumulative_performance.positions[133].sid, 
-            133, 
-            "Portfolio should have one position in 133."
+            zipline.get_positions()[SID]['sid'], 
+            SID, 
+            "Portfolio should have one position in " + str(SID)
         )

zipline/test/test_perf_tracking.py

@@ -9,7 +9,7 @@ import zipline.util as qutil
 import zipline.finance.performance as perf
 import zipline.finance.risk as risk
 import zipline.protocol as zp
-from zipline.finance.trading import TradeSimulationClient
+from zipline.finance.trading import TradeSimulationClient, TradingEnvironment
 class PerformanceTestCase(unittest.TestCase):
     
     def setUp(self):
@@ -17,7 +17,7 @@ class PerformanceTestCase(unittest.TestCase):
         self.benchmark_returns, self.treasury_curves = \
         factory.load_market_data()
         
-        self.trading_environment = risk.TradingEnvironment(
+        self.trading_environment = TradingEnvironment(
             self.benchmark_returns, 
             self.treasury_curves
         )

zipline/test/test_protocol.py

@@ -0,0 +1,130 @@
+"""
+Test the FRAME/UNFRAME functions in the sequence expected from ziplines.    
+"""
+import pytz
+
+from unittest2 import TestCase
+from datetime import datetime, timedelta
+from collections import defaultdict
+
+from nose.tools import timed
+
+import zipline.test.factory as factory
+import zipline.util as qutil
+import zipline.protocol as zp
+
+from zipline.sources import SpecificEquityTrades
+
+DEFAULT_TIMEOUT = 5 # seconds
+
+class ProtocolTestCase(TestCase):
+
+    leased_sockets = defaultdict(list)
+
+    def setUp(self):
+        qutil.configure_logging()
+        self.trading_environment = factory.create_trading_environment()
+
+    @timed(DEFAULT_TIMEOUT)
+    def test_trade_feed_protocol(self):
+
+        sid = 133
+        price = [10.0] * 4
+        volume = [100] * 4
+
+        start_date = datetime.strptime("02/15/2012","%m/%d/%Y")
+        one_day_td = timedelta(days=1)
+
+        trades = factory.create_trade_history(
+            sid, 
+            price, 
+            volume, 
+            start_date, 
+            one_day_td, 
+            self.trading_environment
+        )
+
+        for trade in trades:
+            #simulate data source sending frame
+            msg = zp.DATASOURCE_FRAME(zp.namedict(trade))
+            #feed unpacking frame
+            recovered_trade = zp.DATASOURCE_UNFRAME(msg)
+            #feed sending frame
+            feed_msg = zp.FEED_FRAME(recovered_trade)
+            #transform unframing
+            recovered_feed = zp.FEED_UNFRAME(feed_msg)
+            #do a transform
+            trans_msg = zp.TRANSFORM_FRAME('helloworld', 2345.6)
+            #simulate passthrough transform -- passthrough shouldn't even
+            # unpack the msg, just resend.
+
+            passthrough_msg = zp.TRANSFORM_FRAME(zp.TRANSFORM_TYPE.PASSTHROUGH,\
+                    feed_msg)
+
+            #merge unframes transform and passthrough
+            trans_recovered = zp.TRANSFORM_UNFRAME(trans_msg)
+            pt_recovered = zp.TRANSFORM_UNFRAME(passthrough_msg)
+            #simulated merge
+            pt_recovered.PASSTHROUGH.merge(trans_recovered)
+            #frame the merged event
+            merged_msg = zp.MERGE_FRAME(pt_recovered.PASSTHROUGH)
+            #unframe the merge and validate values
+            event = zp.MERGE_UNFRAME(merged_msg)
+
+            #check the transformed value, should only be in event, not trade.
+            self.assertTrue(event.helloworld == 2345.6)
+            event.delete('helloworld')
+
+            self.assertEqual(zp.namedict(trade), event)
+
+    @timed(DEFAULT_TIMEOUT)
+    def test_order_protocol(self):
+        #client places an order
+        now = datetime.utcnow().replace(tzinfo=pytz.utc)
+        order = zp.namedict({
+            'dt':now,
+            'sid':133,
+            'amount':100
+        })
+        order_msg = zp.ORDER_FRAME(order)
+
+        #order datasource receives
+        order = zp.ORDER_UNFRAME(order_msg)
+        self.assertEqual(order.sid, 133)
+        self.assertEqual(order.amount, 100)
+        self.assertEqual(order.dt, now)
+        
+        #order datasource datasource frames the order
+        order_event = zp.namedict({
+            "sid"        : order.sid,
+            "amount"     : order.amount,
+            "dt"         : order.dt,
+            "source_id"  : zp.FINANCE_COMPONENT.ORDER_SOURCE,
+            "type"       : zp.DATASOURCE_TYPE.ORDER
+        })
+
+
+        order_ds_msg = zp.DATASOURCE_FRAME(order_event)
+
+        #transaction transform unframes
+        recovered_order = zp.DATASOURCE_UNFRAME(order_ds_msg)
+
+        self.assertEqual(now, recovered_order.dt)
+
+        #create a transaction from the order
+        txn = zp.namedict({
+            'sid'        : recovered_order.sid,
+            'amount'     : recovered_order.amount,
+            'dt'         : recovered_order.dt,
+            'price'      : 10.0,
+            'commission' : 0.50
+        })
+
+        #frame that transaction
+        txn_msg = zp.TRANSFORM_FRAME(zp.TRANSFORM_TYPE.TRANSACTION, txn)
+
+        #unframe
+        recovered_tx = zp.TRANSFORM_UNFRAME(txn_msg).TRANSACTION
+        self.assertEqual(recovered_tx.sid, 133)
+        self.assertEqual(recovered_tx.amount, 100)
+

zipline/test/test_risk.py

@@ -7,6 +7,8 @@ import zipline.finance.risk as risk
 import zipline.test.factory as factory
 import zipline.util as qutil
 
+from zipline.finance.trading import TradingEnvironment
+
 class Risk(unittest.TestCase):
     
     def setUp(self):
@@ -17,7 +19,7 @@ class Risk(unittest.TestCase):
         self.benchmark_returns, self.treasury_curves = \
         factory.load_market_data()
         
-        self.trading_calendar = risk.TradingEnvironment(
+        self.trading_env = TradingEnvironment(
             self.benchmark_returns, 
             self.treasury_curves
         )
@@ -27,9 +29,9 @@ class Risk(unittest.TestCase):
         self.tradingday = datetime.timedelta(hours=6, minutes=30)
         self.dt = datetime.datetime.utcnow()
         
-        self.algo_returns_06 = factory.create_returns_from_list(RETURNS, start_date, self.trading_calendar)
+        self.algo_returns_06 = factory.create_returns_from_list(RETURNS, start_date, self.trading_env)
         
-        self.metrics_06 = risk.RiskReport(self.algo_returns_06, self.trading_calendar)
+        self.metrics_06 = risk.RiskReport(self.algo_returns_06, self.trading_env)
         
     def tearDown(self):
         return
@@ -37,21 +39,21 @@ class Risk(unittest.TestCase):
     def test_factory(self):
         returns = [0.1] * 100
         start_date = datetime.datetime(year=2006, month=1, day=1, tzinfo=pytz.utc) 
-        r_objects = factory.create_returns_from_list(returns, start_date, self.trading_calendar)
+        r_objects = factory.create_returns_from_list(returns, start_date, self.trading_env)
         self.assertTrue(r_objects[-1].date <= datetime.datetime(year=2006, month=12, day=31, tzinfo=pytz.utc))
            
     def test_drawdown(self):
         start_date = datetime.datetime(year=2006, month=1, day=1)
-        returns = factory.create_returns_from_list([1.0,-0.5,0.8,.17,1.0,-0.1,-0.45], start_date, self.trading_calendar)
+        returns = factory.create_returns_from_list([1.0,-0.5,0.8,.17,1.0,-0.1,-0.45], start_date, self.trading_env)
         #200, 100, 180, 210.6, 421.2, 379.8, 208.494
-        metrics = risk.RiskMetrics(returns[0].date, returns[-1].date, returns, self.trading_calendar)
+        metrics = risk.RiskMetrics(returns[0].date, returns[-1].date, returns, self.trading_env)
         self.assertEqual(metrics.max_drawdown, 0.505)
     
     def test_benchmark_returns_06(self):
         start_date = datetime.datetime(year=2006, month=1, day=1)
         end_date = datetime.datetime(year=2006, month=12, day=31)
-        returns = factory.create_returns_from_range(start_date, end_date, self.trading_calendar)
-        metrics = risk.RiskReport(returns, self.trading_calendar)
+        returns = factory.create_returns_from_range(start_date, end_date, self.trading_env)
+        metrics = risk.RiskReport(returns, self.trading_env)
         self.assertEqual([round(x.benchmark_period_returns, 4)  for x in metrics.month_periods],    
         [0.0255,0.0005,0.0111,0.0122,-0.0309,0.0001,0.0051,0.0213,0.0246,0.0315,0.0165,0.0126])
         self.assertEqual([round(x.benchmark_period_returns, 4) for x in metrics.three_month_periods],
@@ -63,16 +65,16 @@ class Risk(unittest.TestCase):
     def test_trading_days_06(self):
         start_date = datetime.datetime(year=2006, month=1, day=1)
         end_date = datetime.datetime(year=2006, month=12, day=31)
-        returns = factory.create_returns_from_range(start_date, end_date, self.trading_calendar)
-        metrics = risk.RiskReport(returns, self.trading_calendar)
+        returns = factory.create_returns_from_range(start_date, end_date, self.trading_env)
+        metrics = risk.RiskReport(returns, self.trading_env)
         self.assertEqual([x.trading_days for x in metrics.year_periods],[251])
         self.assertEqual([x.trading_days for x in metrics.month_periods],[20,19,23,19,22,22,20,23,20,22,21,20])
 
     def test_benchmark_volatility_06(self):
         start_date = datetime.datetime(year=2006, month=1, day=1)
         end_date = datetime.datetime(year=2006, month=12, day=31)
-        returns = factory.create_returns_from_range(start_date, end_date, self.trading_calendar)
-        metrics = risk.RiskReport(returns, self.trading_calendar)
+        returns = factory.create_returns_from_range(start_date, end_date, self.trading_env)
+        metrics = risk.RiskReport(returns, self.trading_env)
         self.assertEqual([round(x.benchmark_volatility, 3) for x in metrics.month_periods],
         [0.031,0.026,0.024,0.025,0.037,0.047,0.039,0.022,0.023,0.021,0.025,0.019])
         self.assertEqual([round(x.benchmark_volatility, 3) for x in metrics.three_month_periods],
@@ -133,8 +135,8 @@ class Risk(unittest.TestCase):
     def test_benchmark_returns_08(self):
         start_date = datetime.datetime(year=2008, month=1, day=1)
         end_date = datetime.datetime(year=2008, month=12, day=31)
-        returns = factory.create_returns_from_range(start_date, end_date, self.trading_calendar)
-        metrics = risk.RiskReport(returns, self.trading_calendar)
+        returns = factory.create_returns_from_range(start_date, end_date, self.trading_env)
+        metrics = risk.RiskReport(returns, self.trading_env)
         self.assertEqual([round(x.benchmark_period_returns, 3)  for x in metrics.month_periods],    
         [-0.061,-0.035,-0.006,0.048,0.011,-0.086,-0.01,0.012,-0.091,-0.169,-0.075,0.008])
         self.assertEqual([round(x.benchmark_period_returns, 3) for x in metrics.three_month_periods],
@@ -146,16 +148,16 @@ class Risk(unittest.TestCase):
     def test_trading_days_08(self):
         start_date = datetime.datetime(year=2008, month=1, day=1)
         end_date = datetime.datetime(year=2008, month=12, day=31)
-        returns = factory.create_returns_from_range(start_date, end_date, self.trading_calendar)
-        metrics = risk.RiskReport(returns, self.trading_calendar)
+        returns = factory.create_returns_from_range(start_date, end_date, self.trading_env)
+        metrics = risk.RiskReport(returns, self.trading_env)
         self.assertEqual([x.trading_days for x in metrics.year_periods],[253])
         self.assertEqual([x.trading_days for x in metrics.month_periods],[21,20,20,22,21,21,22,21,21,23,19,22])
         
     def test_benchmark_volatility_08(self):
         start_date = datetime.datetime(year=2008, month=1, day=1)
         end_date = datetime.datetime(year=2008, month=12, day=31)
-        returns = factory.create_returns_from_range(start_date, end_date, self.trading_calendar)
-        metrics = risk.RiskReport(returns, self.trading_calendar)
+        returns = factory.create_returns_from_range(start_date, end_date, self.trading_env)
+        metrics = risk.RiskReport(returns, self.trading_env)
         self.assertEqual([round(x.benchmark_volatility, 3) for x in metrics.month_periods],
         [0.07,0.058,0.082,0.054,0.041,0.057,0.068,0.06,0.157,0.244,0.195,0.145])
         self.assertEqual([round(x.benchmark_volatility, 3) for x in metrics.three_month_periods],
@@ -168,8 +170,8 @@ class Risk(unittest.TestCase):
     def test_treasury_returns_06(self):
         start_date = datetime.datetime(year=2006, month=1, day=1)
         end_date = datetime.datetime(year=2006, month=12, day=31)
-        returns = factory.create_returns_from_range(start_date, end_date, self.trading_calendar)
-        metrics = risk.RiskReport(returns, self.trading_calendar)
+        returns = factory.create_returns_from_range(start_date, end_date, self.trading_env)
+        metrics = risk.RiskReport(returns, self.trading_env)
         self.assertEqual([round(x.treasury_period_return, 4) for x in metrics.month_periods],
         [0.0037,0.0034,0.0039,0.0038,0.0040,0.0037,0.0043,0.0043,0.0038,0.0044,0.0043,0.0041])
         self.assertEqual([round(x.treasury_period_return, 4) for x in metrics.three_month_periods],
@@ -184,9 +186,9 @@ class Risk(unittest.TestCase):
         
     def test_partial_month(self):
         start_date = datetime.datetime(year=1991, month=1, day=1)
-        returns = factory.create_returns(365 * 5 + 2, start_date, self.trading_calendar) #1992 and 1996 were leap years
+        returns = factory.create_returns(365 * 5 + 2, start_date, self.trading_env) #1992 and 1996 were leap years
         returns = returns[:-10] #truncate the returns series to end mid-month
-        metrics = risk.RiskReport(returns, self.trading_calendar)
+        metrics = risk.RiskReport(returns, self.trading_env)
         total_months = 60
         self.check_metrics(metrics, total_months, start_date)
         
@@ -196,8 +198,8 @@ class Risk(unittest.TestCase):
         else:
             #because we may catch the leap of the last year, and i think this func is [start,end)
             ld = calendar.leapdays(start_date.year, start_date.year + years + 1) 
-        returns = factory.create_returns(365 * years + ld, start_date, self.trading_calendar)
-        metrics = risk.RiskReport(returns, self.trading_calendar)
+        returns = factory.create_returns(365 * years + ld, start_date, self.trading_env)
+        metrics = risk.RiskReport(returns, self.trading_env)
         total_months = years * 12
         self.check_metrics(metrics, total_months, start_date)