diff --git a/zipline/testing/fixtures.py b/zipline/testing/fixtures.py index 76fb6bcb..8ee6378b 100644 --- a/zipline/testing/fixtures.py +++ b/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