Merge pull request #100 from quantopian/bug_524

Bug 524

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')
         

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

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):

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)