diff --git a/tests/test_exception_handling.py b/tests/test_exception_handling.py index ac6339bc..f16111ee 100644 --- a/tests/test_exception_handling.py +++ b/tests/test_exception_handling.py @@ -9,6 +9,8 @@ from zipline.finance.trading import SIMULATION_STYLE from zipline.core.devsimulator import AddressAllocator from zipline.lines import SimulatedTrading from zipline.gens.transform import StatefulTransform +from zipline.gens.tradesimulation import HEARTBEAT_INTERVAL, \ + MAX_HEARTBEAT_INTERVALS from zipline.utils.test_utils import \ drain_zipline, \ @@ -161,18 +163,29 @@ class ExceptionTestCase(TestCase): self.assertEqual(payload['name'],'Timeout') self.assertEqual(payload['message'], 'Call to initialize timed out') -# def test_heartbeat(self): + def test_heartbeat(self): -# self.zipline_test_config['algorithm'] = \ -# TooMuchProcessingAlgorithm( -# self.zipline_test_config['sid'] -# ) -# zipline = SimulatedTrading.create_test_zipline( -# **self.zipline_test_config -# ) -# output, _ = drain_zipline(self, zipline) -# self.assertEqual(output[-1]['prefix'], 'EXCEPTION') -# payload = output[-1]['payload'] -# self.assertEqual(payload['name'],'Timeout') -# self.assertEqual(payload['message'], 'Too much time spent in handle_data call') + self.zipline_test_config['algorithm'] = \ + TooMuchProcessingAlgorithm( + self.zipline_test_config['sid'] + ) + zipline = SimulatedTrading.create_test_zipline( + **self.zipline_test_config + ) + output, _ = drain_zipline(self, zipline) + + # There should be a message for each hearbeat, plus a message + # for the final timeout. + assert len(output) == MAX_HEARTBEAT_INTERVALS + 1 + + # Assert that everything but the last message is a heartbeat log. + for message in output[0:-1]: + assert message['prefix'] == 'LOG' + assert message['payload']['func_name'] == 'log_heartbeats' + + # Assert that the last message is a timeout exception. + self.assertEqual(output[-1]['prefix'], 'EXCEPTION') + payload = output[-1]['payload'] + self.assertEqual(payload['name'],'Timeout') + self.assertEqual(payload['message'], 'Too much time spent in handle_data call') diff --git a/zipline/gens/tradesimulation.py b/zipline/gens/tradesimulation.py index 811f99ac..02a9e03c 100644 --- a/zipline/gens/tradesimulation.py +++ b/zipline/gens/tradesimulation.py @@ -61,7 +61,6 @@ class TradeSimulationClient(object): self.sids = algo.get_sid_filter() self.environment = environment self.style = sim_style - self.algo_sim = None self.warmup_start = self.environment.prior_day_open self.algo_start = self.environment.first_open @@ -150,8 +149,10 @@ class AlgorithmSimulator(object): # Handler for heartbeats during calls to handle_data. def log_heartbeats(beat_count, stackframe): t = beat_count * HEARTBEAT_INTERVAL - warning = "handle_data has been processing for %i seconds" %t + warning = "handle_data has been processing for %i seconds" %t + # Log the warning to our logs as well as the user's log output. self.algolog.warn(warning) + log.warn(warning) # Context manager that calls log_heartbeats every HEARTBEAT_INTERVAL # seconds, raising an exception after MAX_HEARTBEATS diff --git a/zipline/test_algorithms.py b/zipline/test_algorithms.py index f6bdfd4d..eb883469 100644 --- a/zipline/test_algorithms.py +++ b/zipline/test_algorithms.py @@ -265,7 +265,7 @@ class TooMuchProcessingAlgorithm(): def handle_data(self, data): # Unless we're running on some sort of # supercomputer this will hit timeout. - for i in xrange(100000000): + for i in xrange(1000000000): self.foo = i def get_sid_filter(self): diff --git a/zipline/utils/timeout.py b/zipline/utils/timeout.py index 0310a208..2985222e 100644 --- a/zipline/utils/timeout.py +++ b/zipline/utils/timeout.py @@ -106,6 +106,7 @@ class heartbeat(object): # decorator doesn't have unexpected side-effects later. signal.setitimer(signal.ITIMER_REAL, 0, 0) signal.signal(signal.SIGALRM, signal.SIG_DFL) + self.count = 0 # Return the value of fn if it finished without tripping # an exception. This won't execute if the Timeout or any @@ -115,11 +116,13 @@ class heartbeat(object): def __enter__(self): # Set a timer to call our handler every N seconds. + self.count = 0 signal.signal(signal.SIGALRM, self.handler) signal.setitimer(signal.ITIMER_REAL, self.interval, self.interval) def __exit__(self, type, value, traceback): # Turn off the timer on exit. This will re-raise any exception raised # during execution of the with-block + self.count = 0 signal.setitimer(signal.ITIMER_REAL, 0, 0) signal.signal(signal.SIGALRM, signal.SIG_DFL)