TST: Reconciles tests with asset management system

This commit is contained in:
jfkirk
2015-06-11 11:35:49 -04:00
parent b84ac01cbf
commit a5d1f79a37
11 changed files with 608 additions and 391 deletions
+55 -36
View File
@@ -6,6 +6,7 @@ from unittest import TestCase
from zipline.algorithm import TradingAlgorithm
from zipline.errors import TradingControlViolation
from zipline.sources import SpecificEquityTrades
from zipline.finance.trading import with_environment
from zipline.utils.test_utils import (
setup_logger, teardown_logger, security_list_copy, add_security_data)
from zipline.utils import factory
@@ -16,11 +17,11 @@ LEVERAGED_ETFS = load_from_directory('leveraged_etf_list')
class RestrictedAlgoWithCheck(TradingAlgorithm):
def initialize(self, sid):
def initialize(self, symbol):
self.rl = SecurityListSet(self.get_datetime)
self.set_do_not_order_list(self.rl.leveraged_etf_list)
self.order_count = 0
self.sid = sid
self.sid = self.symbol(symbol)
def handle_data(self, data):
if not self.order_count:
@@ -31,11 +32,11 @@ class RestrictedAlgoWithCheck(TradingAlgorithm):
class RestrictedAlgoWithoutCheck(TradingAlgorithm):
def initialize(self, sid):
def initialize(self, symbol):
self.rl = SecurityListSet(self.get_datetime)
self.set_do_not_order_list(self.rl.leveraged_etf_list)
self.order_count = 0
self.sid = sid
self.sid = self.symbol(symbol)
def handle_data(self, data):
self.order(self.sid, 100)
@@ -43,11 +44,11 @@ class RestrictedAlgoWithoutCheck(TradingAlgorithm):
class IterateRLAlgo(TradingAlgorithm):
def initialize(self, sid):
def initialize(self, symbol):
self.rl = SecurityListSet(self.get_datetime)
self.set_do_not_order_list(self.rl.leveraged_etf_list)
self.order_count = 0
self.sid = sid
self.sid = self.symbol(symbol)
self.found = False
def handle_data(self, data):
@@ -58,12 +59,18 @@ class IterateRLAlgo(TradingAlgorithm):
class SecurityListTestCase(TestCase):
def setUp(self):
@with_environment()
def setUp(self, env=None):
self.extra_knowledge_date = \
datetime(2015, 1, 27, 0, 0, tzinfo=pytz.utc)
self.trading_day_before_first_kd = datetime(
2015, 1, 23, 0, 0, tzinfo=pytz.utc)
env.update_asset_finder(
clear_metadata=True,
identifiers=["BZQ", "URTY", "JFT", "AAPL", "GOOG"]
)
setup_logger(self)
def tearDown(self):
@@ -81,11 +88,12 @@ class SecurityListTestCase(TestCase):
sim_params
)
self.source = SpecificEquityTrades(event_list=trade_history)
algo = IterateRLAlgo(sid='BZQ', sim_params=sim_params)
algo = IterateRLAlgo(symbol='BZQ', sim_params=sim_params)
algo.run(self.source)
self.assertTrue(algo.found)
def test_security_list(self):
@with_environment()
def test_security_list(self, env=None):
# set the knowledge date to the first day of the
# leveraged etf knowledge date.
@@ -94,27 +102,43 @@ class SecurityListTestCase(TestCase):
rl = SecurityListSet(get_datetime)
# assert that a sample from the leveraged list are in restricted
self.assertIn("BZQ", rl.leveraged_etf_list)
self.assertIn("URTY", rl.leveraged_etf_list)
self.assertIn("JFT", rl.leveraged_etf_list)
should_exist = [
asset.sid for asset in
[env.asset_finder.lookup_symbol(
symbol,
as_of_date=self.extra_knowledge_date)
for symbol in ["BZQ", "URTY", "JFT"]]
]
for sid in should_exist:
self.assertIn(sid, rl.leveraged_etf_list)
# assert that a sample of allowed stocks are not in restricted
# AAPL
self.assertNotIn("AAPL", rl.leveraged_etf_list)
# GOOG
self.assertNotIn("GOOG", rl.leveraged_etf_list)
shouldnt_exist = [
asset.sid for asset in
[env.asset_finder.lookup_symbol(
symbol,
as_of_date=self.extra_knowledge_date)
for symbol in ["AAPL", "GOOG"]]
]
for sid in shouldnt_exist:
self.assertNotIn(sid, rl.leveraged_etf_list)
def test_security_add(self):
@with_environment()
def test_security_add(self, env=None):
def get_datetime():
return datetime(2015, 1, 27, tzinfo=pytz.utc)
with security_list_copy():
add_security_data(['AAPL', 'GOOG'], [])
rl = SecurityListSet(get_datetime)
self.assertIn("AAPL", rl.leveraged_etf_list)
self.assertIn("GOOG", rl.leveraged_etf_list)
self.assertIn("BZQ", rl.leveraged_etf_list)
self.assertIn("URTY", rl.leveraged_etf_list)
should_exist = [
asset.sid for asset in
[env.asset_finder.lookup_symbol(
symbol,
as_of_date=self.extra_knowledge_date
) for symbol in ["AAPL", "GOOG", "BZQ", "URTY"]]
]
for sid in should_exist:
self.assertIn(sid, rl.leveraged_etf_list)
def test_security_add_delete(self):
with security_list_copy():
@@ -138,7 +162,7 @@ class SecurityListTestCase(TestCase):
)
self.source = SpecificEquityTrades(event_list=trade_history)
algo = RestrictedAlgoWithCheck(sid='BZQ', sim_params=sim_params)
algo = RestrictedAlgoWithCheck(symbol='BZQ', sim_params=sim_params)
algo.run(self.source)
def test_algo_without_rl_violation(self):
@@ -153,7 +177,7 @@ class SecurityListTestCase(TestCase):
sim_params
)
self.source = SpecificEquityTrades(event_list=trade_history)
algo = RestrictedAlgoWithoutCheck(sid='AAPL', sim_params=sim_params)
algo = RestrictedAlgoWithoutCheck(symbol='AAPL', sim_params=sim_params)
algo.run(self.source)
def test_algo_with_rl_violation(self):
@@ -169,10 +193,7 @@ class SecurityListTestCase(TestCase):
)
self.source = SpecificEquityTrades(event_list=trade_history)
self.df_source, self.df = \
factory.create_test_df_source(sim_params)
algo = RestrictedAlgoWithoutCheck(sid='BZQ', sim_params=sim_params)
algo = RestrictedAlgoWithoutCheck(symbol='BZQ', sim_params=sim_params)
with self.assertRaises(TradingControlViolation) as ctx:
algo.run(self.source)
@@ -189,10 +210,7 @@ class SecurityListTestCase(TestCase):
)
self.source = SpecificEquityTrades(event_list=trade_history)
self.df_source, self.df = \
factory.create_test_df_source(sim_params)
algo = RestrictedAlgoWithoutCheck(sid='JFT', sim_params=sim_params)
algo = RestrictedAlgoWithoutCheck(symbol='JFT', sim_params=sim_params)
with self.assertRaises(TradingControlViolation) as ctx:
algo.run(self.source)
@@ -211,7 +229,7 @@ class SecurityListTestCase(TestCase):
sim_params
)
self.source = SpecificEquityTrades(event_list=trade_history)
algo = RestrictedAlgoWithoutCheck(sid='BZQ', sim_params=sim_params)
algo = RestrictedAlgoWithoutCheck(symbol='BZQ', sim_params=sim_params)
with self.assertRaises(TradingControlViolation) as ctx:
algo.run(self.source)
@@ -238,7 +256,7 @@ class SecurityListTestCase(TestCase):
)
self.source = SpecificEquityTrades(event_list=trade_history)
algo = RestrictedAlgoWithoutCheck(
sid='BZQ', sim_params=sim_params)
symbol='BZQ', sim_params=sim_params)
with self.assertRaises(TradingControlViolation) as ctx:
algo.run(self.source)
@@ -261,7 +279,8 @@ class SecurityListTestCase(TestCase):
)
self.source = SpecificEquityTrades(event_list=trade_history)
algo = RestrictedAlgoWithoutCheck(
sid='BZQ', sim_params=sim_params)
symbol='BZQ', sim_params=sim_params
)
algo.run(self.source)
def test_algo_with_rl_violation_after_add(self):
@@ -278,7 +297,7 @@ class SecurityListTestCase(TestCase):
)
self.source = SpecificEquityTrades(event_list=trade_history)
algo = RestrictedAlgoWithoutCheck(
sid='AAPL', sim_params=sim_params)
symbol='AAPL', sim_params=sim_params)
with self.assertRaises(TradingControlViolation) as ctx:
algo.run(self.source)