mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-13 17:42:42 +08:00
TST: Update test_algorithm.py to incorporate TradingEnvironment.write_data
This commit is contained in:
+44
-17
@@ -94,6 +94,7 @@ from zipline.sources import (SpecificEquityTrades,
|
||||
from zipline.assets import Equity
|
||||
|
||||
from zipline.finance.execution import LimitOrder
|
||||
from zipline.finance import trading
|
||||
from zipline.finance.trading import SimulationParameters
|
||||
from zipline.utils.api_support import set_algo_instance
|
||||
from zipline.utils.events import DateRuleFactory, TimeRuleFactory
|
||||
@@ -109,6 +110,7 @@ _multiprocess_can_split_ = False
|
||||
class TestRecordAlgorithm(TestCase):
|
||||
def setUp(self):
|
||||
self.sim_params = factory.create_simulation_parameters(num_days=4)
|
||||
trading.environment.write_data(equities_identifiers=[133])
|
||||
trade_history = factory.create_trade_history(
|
||||
133,
|
||||
[10.0, 10.0, 11.0, 11.0],
|
||||
@@ -145,6 +147,7 @@ class TestMiscellaneousAPI(TestCase):
|
||||
data_frequency='minute',
|
||||
emission_rate='minute',
|
||||
)
|
||||
trading.environment.write_data(equities_identifiers=sids)
|
||||
self.source = factory.create_minutely_trade_source(
|
||||
sids,
|
||||
sim_params=self.sim_params,
|
||||
@@ -321,6 +324,8 @@ class TestMiscellaneousAPI(TestCase):
|
||||
self.assertIs(composer, zipline.utils.events.ComposedRule.lazy_and)
|
||||
|
||||
def test_asset_lookup(self):
|
||||
|
||||
trading.environment = trading.TradingEnvironment()
|
||||
metadata = {0: {'symbol': 'PLAY',
|
||||
'asset_type': 'equity',
|
||||
'start_date': '2002-01-01',
|
||||
@@ -364,7 +369,7 @@ class TestMiscellaneousAPI(TestCase):
|
||||
def test_future_chain(self):
|
||||
""" Tests the future_chain API function.
|
||||
"""
|
||||
|
||||
trading.environment = trading.TradingEnvironment()
|
||||
metadata = {
|
||||
0: {
|
||||
'symbol': 'CLG06',
|
||||
@@ -396,7 +401,8 @@ class TestMiscellaneousAPI(TestCase):
|
||||
'expiration_date': pd.Timestamp('2006-10-20', tz='UTC')}
|
||||
}
|
||||
|
||||
algo = TradingAlgorithm(asset_metadata=metadata)
|
||||
trading.environment.write_data(futures_data=metadata)
|
||||
algo = TradingAlgorithm()
|
||||
algo.datetime = pd.Timestamp('2006-12-01', tz='UTC')
|
||||
|
||||
# Check that the fields of the FutureChain object are set correctly
|
||||
@@ -483,6 +489,7 @@ class TestTransformAlgorithm(TestCase):
|
||||
def setUp(self):
|
||||
setup_logger(self)
|
||||
self.sim_params = factory.create_simulation_parameters(num_days=4)
|
||||
trading.environment.write_data(equities_identifiers=[0, 1, 133])
|
||||
|
||||
trade_history = factory.create_trade_history(
|
||||
133,
|
||||
@@ -591,9 +598,9 @@ class TestTransformAlgorithm(TestCase):
|
||||
(TestTargetValueAlgorithm,),
|
||||
])
|
||||
def test_order_methods_for_future(self, algo_class):
|
||||
metadata = {0: {'asset_type': 'future',
|
||||
# Use sid not already in test database.
|
||||
metadata = {3: {'asset_type': 'future',
|
||||
'contract_multiplier': 10}}
|
||||
|
||||
algo = algo_class(
|
||||
sim_params=self.sim_params,
|
||||
asset_metadata=metadata
|
||||
@@ -640,6 +647,7 @@ class TestPositions(TestCase):
|
||||
def setUp(self):
|
||||
setup_logger(self)
|
||||
self.sim_params = factory.create_simulation_parameters(num_days=4)
|
||||
trading.environment.write_data(equities_identifiers=[0, 1, 133])
|
||||
|
||||
trade_history = factory.create_trade_history(
|
||||
1,
|
||||
@@ -684,9 +692,10 @@ class TestPositions(TestCase):
|
||||
class TestAlgoScript(TestCase):
|
||||
def setUp(self):
|
||||
days = 251
|
||||
# Note that create_simulation_parameters creates a new TradingEnvironment
|
||||
self.sim_params = factory.create_simulation_parameters(num_days=days)
|
||||
setup_logger(self)
|
||||
|
||||
trading.environment.write_data(equities_identifiers=[0, 1, 133])
|
||||
trade_history = factory.create_trade_history(
|
||||
133,
|
||||
[10.0] * days,
|
||||
@@ -727,7 +736,8 @@ class TestAlgoScript(TestCase):
|
||||
|
||||
def test_api_get_environment(self):
|
||||
platform = 'zipline'
|
||||
metadata = {0: {'symbol': 'TEST',
|
||||
# Use sid not already in test database.
|
||||
metadata = {3: {'symbol': 'TEST',
|
||||
'asset_type': 'equity'}}
|
||||
algo = TradingAlgorithm(script=api_get_environment_algo,
|
||||
asset_metadata=metadata,
|
||||
@@ -736,7 +746,8 @@ class TestAlgoScript(TestCase):
|
||||
self.assertEqual(algo.environment, platform)
|
||||
|
||||
def test_api_symbol(self):
|
||||
metadata = {0: {'symbol': 'TEST',
|
||||
# Use sid not already in test database.
|
||||
metadata = {3: {'symbol': 'TEST',
|
||||
'asset_type': 'equity'}}
|
||||
algo = TradingAlgorithm(script=api_symbol_algo,
|
||||
asset_metadata=metadata)
|
||||
@@ -983,6 +994,15 @@ def handle_data(context, data):
|
||||
|
||||
|
||||
class TestHistory(TestCase):
|
||||
|
||||
def setUp(self):
|
||||
setup_logger(self)
|
||||
trading.environment = trading.TradingEnvironment()
|
||||
trading.environment.write_data(equities_identifiers=[0, 1])
|
||||
|
||||
def tearDown(self):
|
||||
teardown_logger(self)
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls._start = pd.Timestamp('1991-01-01', tz='UTC')
|
||||
@@ -1044,6 +1064,14 @@ def handle_data(context, data):
|
||||
|
||||
class TestGetDatetime(TestCase):
|
||||
|
||||
def setUp(self):
|
||||
setup_logger(self)
|
||||
trading.environment = trading.TradingEnvironment()
|
||||
trading.environment.write_data(equities_identifiers=[0, 1])
|
||||
|
||||
def tearDown(self):
|
||||
teardown_logger(self)
|
||||
|
||||
@parameterized.expand(
|
||||
[
|
||||
('default', None,),
|
||||
@@ -1087,7 +1115,6 @@ class TestGetDatetime(TestCase):
|
||||
algo = TradingAlgorithm(
|
||||
script=algo,
|
||||
sim_params=sim_params,
|
||||
identifiers=[1]
|
||||
)
|
||||
algo.run(source)
|
||||
self.assertFalse(algo.first_bar)
|
||||
@@ -1098,6 +1125,7 @@ class TestTradingControls(TestCase):
|
||||
def setUp(self):
|
||||
self.sim_params = factory.create_simulation_parameters(num_days=4)
|
||||
self.sid = 133
|
||||
trading.environment.write_data(equities_identifiers=[self.sid])
|
||||
self.trade_history = factory.create_trade_history(
|
||||
self.sid,
|
||||
[10.0, 10.0, 11.0, 11.0],
|
||||
@@ -1351,31 +1379,28 @@ class TestTradingControls(TestCase):
|
||||
df_source, _ = factory.create_test_df_source(self.sim_params)
|
||||
metadata = {0: {'start_date': '1990-01-01',
|
||||
'end_date': '2020-01-01'}}
|
||||
asset_finder = AssetFinder()
|
||||
algo = SetAssetDateBoundsAlgorithm(
|
||||
asset_finder=asset_finder,
|
||||
asset_metadata=metadata,
|
||||
sim_params=self.sim_params,)
|
||||
algo.run(df_source)
|
||||
|
||||
# Run the algorithm with a sid that has already ended
|
||||
trading.environment = trading.TradingEnvironment()
|
||||
df_source, _ = factory.create_test_df_source(self.sim_params)
|
||||
metadata = {0: {'start_date': '1989-01-01',
|
||||
'end_date': '1990-01-01'}}
|
||||
asset_finder = AssetFinder()
|
||||
algo = SetAssetDateBoundsAlgorithm(
|
||||
asset_finder=asset_finder,
|
||||
asset_metadata=metadata,
|
||||
sim_params=self.sim_params,)
|
||||
with self.assertRaises(TradingControlViolation):
|
||||
algo.run(df_source)
|
||||
|
||||
# Run the algorithm with a sid that has not started
|
||||
trading.environment = trading.TradingEnvironment()
|
||||
df_source, _ = factory.create_test_df_source(self.sim_params)
|
||||
metadata = {0: {'start_date': '2020-01-01',
|
||||
'end_date': '2021-01-01'}}
|
||||
algo = SetAssetDateBoundsAlgorithm(
|
||||
asset_finder=asset_finder,
|
||||
asset_metadata=metadata,
|
||||
sim_params=self.sim_params,)
|
||||
with self.assertRaises(TradingControlViolation):
|
||||
@@ -1387,6 +1412,7 @@ class TestAccountControls(TestCase):
|
||||
def setUp(self):
|
||||
self.sim_params = factory.create_simulation_parameters(num_days=4)
|
||||
self.sidint = 133
|
||||
trading.environment.write_data(equities_identifiers=[self.sidint])
|
||||
self.trade_history = factory.create_trade_history(
|
||||
self.sidint,
|
||||
[10.0, 10.0, 11.0, 11.0],
|
||||
@@ -1436,6 +1462,7 @@ class TestAccountControls(TestCase):
|
||||
class TestClosePosAlgo(TestCase):
|
||||
|
||||
def setUp(self):
|
||||
trading.environment = trading.TradingEnvironment()
|
||||
self.days = TradingEnvironment().trading_days
|
||||
self.index = [self.days[0], self.days[1], self.days[2]]
|
||||
self.panel = pd.Panel({1: pd.DataFrame({
|
||||
@@ -1473,9 +1500,9 @@ class TestClosePosAlgo(TestCase):
|
||||
metadata = {1: {'symbol': 'TEST',
|
||||
'asset_type': 'future',
|
||||
}}
|
||||
trading.environment.write_data(futures_data=metadata)
|
||||
self.algo = TestAlgorithm(sid=1, amount=1, order_count=1,
|
||||
instant_fill=True, commission=PerShare(0),
|
||||
asset_metadata=metadata)
|
||||
instant_fill=True, commission=PerShare(0),)
|
||||
self.data = DataPanelSource(self.panel)
|
||||
|
||||
# Check results
|
||||
@@ -1490,9 +1517,9 @@ class TestClosePosAlgo(TestCase):
|
||||
'asset_type': 'future',
|
||||
'notice_date': self.days[3],
|
||||
'expiration_date': self.days[4]}}
|
||||
trading.environment.write_data(futures_data=metadata)
|
||||
self.algo = TestAlgorithm(sid=1, amount=1, order_count=1,
|
||||
instant_fill=True, commission=PerShare(0),
|
||||
asset_metadata=metadata)
|
||||
instant_fill=True, commission=PerShare(0),)
|
||||
self.data = DataPanelSource(self.no_close_panel)
|
||||
|
||||
# Check results
|
||||
|
||||
Reference in New Issue
Block a user