mirror of
https://github.com/wassname/catalyst.git
synced 2026-06-30 03:11:14 +08:00
scottsanderson
102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
from datetime import timedelta
|
|
from collections import defaultdict
|
|
from unittest2 import TestCase
|
|
|
|
from zipline.utils.test_utils import setup_logger, teardown_logger
|
|
|
|
import zipline.utils.factory as factory
|
|
|
|
from zipline.gens.tradegens import SpecificEquityTrades
|
|
from zipline.gens.transform import StatefulTransform
|
|
from zipline.gens.vwap import VWAP
|
|
from zipline.gens.mavg import MovingAverage
|
|
from zipline.gens.returns import Returns
|
|
from zipline.lines import SimulatedTrading
|
|
from zipline.core.devsimulator import AddressAllocator
|
|
|
|
allocator = AddressAllocator(1000)
|
|
|
|
class FinanceTransformsTestCase(TestCase):
|
|
|
|
def setUp(self):
|
|
self.trading_environment = factory.create_trading_environment()
|
|
setup_logger(self)
|
|
|
|
trade_history = factory.create_trade_history(
|
|
133,
|
|
[10.0, 10.0, 11.0, 11.0],
|
|
[100, 100, 100, 300],
|
|
timedelta(days=1),
|
|
self.trading_environment
|
|
)
|
|
self.source = SpecificEquityTrades(event_list=trade_history)
|
|
|
|
def tearDown(self):
|
|
self.log_handler.pop_application()
|
|
|
|
def test_vwap(self):
|
|
|
|
vwap = StatefulTransform(
|
|
VWAP,
|
|
market_aware = False,
|
|
delta = timedelta(days = 2)
|
|
)
|
|
transformed = list(vwap.transform(self.source))
|
|
|
|
# Output values
|
|
tnfm_vals = [message.tnfm_value for message in transformed]
|
|
# "Hand calculated" values.
|
|
expected = [(10.0 * 100) / 100.0,
|
|
((10.0 * 100) + (10.0 * 100)) / (200.0),
|
|
# We should drop the first event here.
|
|
((10.0 * 100) + (11.0 * 100)) / (200.0),
|
|
# We should drop the second event here.
|
|
((11.0 * 100) + (11.0 * 300)) / (400.0)]
|
|
|
|
# Output should match the expected.
|
|
assert tnfm_vals == expected
|
|
|
|
def test_returns(self):
|
|
trade_history = factory.create_trade_history(
|
|
133,
|
|
[10.0, 10.0, 10.0, 11.0],
|
|
[100, 100, 100, 300],
|
|
timedelta(days=1),
|
|
)
|
|
returns = StatefulTransform(
|
|
Returns
|
|
for trade in trade_history:
|
|
returns.update(trade)
|
|
self.assertEqual(returns.returns, .1)
|
|
|
|
def test_moving_average(self):
|
|
|
|
mavg = StatefulTransform(
|
|
MovingAverage,
|
|
market_aware = False,
|
|
fields = ['price', 'volume'],
|
|
delta = timedelta(days = 2),
|
|
)
|
|
transformed = list(mavg.transform(self.source))
|
|
# Output values.
|
|
tnfm_prices = [message.tnfm_value.price for message in transformed]
|
|
tnfm_volumes = [message.tnfm_value.volume for message in transformed]
|
|
|
|
# "Hand-calculated" values
|
|
expected_prices = [((10.0) / 1.0),
|
|
((10.0 + 10.0) / 2.0),
|
|
# First event should get dropped here.
|
|
((10.0 + 11.0) / 2.0),
|
|
# Second event should get dropped here.
|
|
((11.0 + 11.0) / 2.0)]
|
|
|
|
expected_volumes = [((100.0) / 1.0),
|
|
((100.0 + 100.0) / 2.0),
|
|
# First event should get dropped here.
|
|
((100.0 + 100.0) / 2.0),
|
|
# Second event should get dropped here.
|
|
((100.0 + 300.0) / 2.0)]
|
|
|
|
assert tnfm_prices == expected_prices
|
|
assert tnfm_volumes == expected_volumes
|