TST: Adds WithLogger and WithTradingEnvironment

zipline/testing/fixtures.py

@@ -1,11 +1,13 @@
 from unittest import TestCase
 
 from contextlib2 import ExitStack
+from logbook import NullHandler
 import pandas as pd
 from six import with_metaclass
 
 from .core import tmp_asset_finder
-from ..utils import tradingcalendar
+from ..finance.trading import TradingEnvironment
+from ..utils import tradingcalendar, factory
 from ..utils.final import FinalMeta, final
 
 
@@ -71,6 +73,11 @@ class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)):
     def add_class_callback(cls, callback):
         """
         Register a callback to be executed during tearDownClass.
+
+        Parameters
+        ----------
+        callback : callable
+            The callback to invoke at the end of the test suite.
         """
         return cls._class_teardown_stack.callback(callback)
 
@@ -108,30 +115,147 @@ class ZiplineTestCase(with_metaclass(FinalMeta, TestCase)):
     def add_instance_callback(self, callback):
         """
         Register a callback to be executed during tearDown.
+
+        Parameters
+        ----------
+        callback : callable
+            The callback to invoke at the end of each test.
         """
         return self._instance_teardown_stack.callback(callback)
 
 
+class WithLogger(object):
+    """ZiplineTestCase mixing providing cls.log_handler as an instance-level
+    fixture.
+
+    after init_instance_fixtures has been called `self.log_handler` will be a
+    new ``NullHandler``.
+
+    This behaviour may be overridden by defining a ``make_log_handler`` class
+    method which returns a new handler object.
+    """
+    make_log_handler = NullHandler
+
+    @classmethod
+    def init_class_fixtures(cls):
+        super(WithLogger, cls).init_instance_fixtures()
+
+        cls.log_handler = cls.enter_class_context(
+            cls.make_log_handler().applicationbound(),
+        )
+
+
 class WithAssetFinder(object):
     """
     ZiplineTestCase mixin providing cls.asset_finder as a class-level fixture.
 
     After init_class_fixtures has been called, `cls.asset_finder` is populated
-    with an AssetFinder. The default finder just a `tmp_asset_finder`
+    with an AssetFinder. The default finder just a `tmp_asset_finder` whose
+    arguments come from:
+
+      equities=cls.make_equities_info(),
+      futures=cls.make_futures_info(),
+      exchanges=cls.make_exchanges_info(),
+      root_symbols=cls.make_root_symbols_info(),
+
+    Each of these fields may be overridden with a function named
+    ``make_*_info`` which returnsa dataframe to write.
 
     This behavior can be altered by overriding `make_asset_finder` as a class
     method.
     """
+    @classmethod
+    def _make_info(cls):
+        return None
+
+    make_equities_info = _make_info
+    make_futures_info = _make_info
+    make_exchanges_info = _make_info
+    make_root_symbols_info = _make_info
+
+    del _make_info
+
     @classmethod
     def make_asset_finder(cls):
-        return cls.enter_class_context(tmp_asset_finder())
+        return cls.enter_class_context(tmp_asset_finder(
+            equities=cls.equities_info,
+            futures=cls.futures_info,
+            exchanges=cls.exchanges_info,
+            root_symbols=cls.root_symbols_info,
+        ))
 
     @classmethod
     def init_class_fixtures(cls):
         super(WithAssetFinder, cls).init_class_fixtures()
+        cls.equities_info = cls.make_equities_info()
+        cls.futures_info = cls.make_futures_info()
+        cls.exchanges_info = cls.make_exchanges_info()
+        cls.root_symbols_info = cls.make_root_symbols_info()
         cls.asset_finder = cls.make_asset_finder()
 
 
+class WithTradingEnvironment(WithAssetFinder):
+    """
+    ZiplineTestCase mixin providing cls.env as a class-level fixture.
+
+    After ``init_class_fixtures`` has been called, `cls.env` is populated
+    with a trading environment whose `asset_finder` is the result of
+    `cls.make_asset_finder`.
+
+    The ``load`` function may be provided by overriding the
+    ``make_load_function`` class method.
+
+    This behavior can be altered by overriding `make_trading_environment` as a
+    class method.
+    """
+    @classmethod
+    def make_load_function(cls):
+        return None
+
+    @classmethod
+    def make_trading_environment(cls):
+        return TradingEnvironment(
+            load=cls.make_load_function(),
+            asset_db_path=cls.asset_finder.engine,
+        )
+
+    @classmethod
+    def init_class_fixtures(cls):
+        super(WithTradingEnvironment, cls).init_class_fixtures()
+        cls.env = cls.make_trading_environment()
+
+
+class WithSimParams(WithTradingEnvironment):
+    """
+    ZiplineTestCase mixin providing cls.sim_params as a class level fixture.
+
+    The arguments used to construct the trading environment may be overridded
+    by putting ``SIM_PARAMS_{argname}`` in the class dict except for the
+    trading environment which is overridden with the mechanisms provided by
+    ``WithTradingEnvironment``.
+    """
+    SIM_PARAMS_YEAR = None
+    SIM_PARAMS_START = pd.Timestamp('2006-01-01')
+    SIM_PARAMS_END = pd.Timestamp('2006-12-31')
+    SIM_PARAMS_CAPITAL_BASE = float("1.0e5")
+    SIM_PARAMS_NUM_DAYS = None
+    SIM_PARAMS_DATA_FREQUENCY = 'daily'
+    SIM_PARAMS_EMISSION_RATE = 'daily'
+
+    @classmethod
+    def init_class_fixtures(cls):
+        super(WithSimParams, cls).init_class_fixtures()
+        cls.sim_params = factory.create_simulation_parameters(
+            year=cls.SIM_PARAMS_YEAR,
+            start=cls.SIM_PARAMS_START,
+            end=cls.SIM_PARAMS_END,
+            capital_base=cls.SIM_PARAMS_CAPITAL_BASE,
+            data_frequency=cls.SIM_PARAMS_DATA_FREQUENCY,
+            emission_rate=cls.SIM_PARAMS_EMISSION_RATE,
+            env=cls.env,
+        )
+
+
 class WithNYSETradingDays(object):
     """
     ZiplineTestCase mixin providing cls.trading_days as a class-level fixture.
@@ -144,7 +268,7 @@ class WithNYSETradingDays(object):
     The default value of TRADING_DAY_COUNT is 126 (half a trading-year).
     Inheritors can override TRADING_DAY_COUNT to request more or less data.
     """
-    DATA_MAX_DAY = pd.Timestamp('2015-01-02')
+    DATA_MAX_DAY = pd.Timestamp('2016-01-04')
     TRADING_DAY_COUNT = 126
 
     @classmethod