added tests for exception messages in datasources and transforms

tests/test_exception_handling.py

@@ -7,12 +7,15 @@ from zipline.test_algorithms import ExceptionAlgorithm, DivByZeroAlgorithm
 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.utils.test_utils import \
         drain_zipline, \
         check, \
         setup_logger, \
-        teardown_logger
+        teardown_logger, \
+        ExceptionSource, \
+        ExceptionTransform
 
 DEFAULT_TIMEOUT = 15 # seconds
 EXTENDED_TIMEOUT = 90
@@ -36,6 +39,39 @@ class ExceptionTestCase(TestCase):
         self.ctx.term()
         teardown_logger(self)
 
+    def test_datasource_exception(self):
+        self.zipline_test_config['trade_source'] = ExceptionSource()
+        zipline = SimulatedTrading.create_test_zipline(
+            **self.zipline_test_config
+        )
+        output, _ = drain_zipline(self, zipline)
+        assert len(output) == 1
+        assert output[0]['prefix'] == 'EXCEPTION'
+        message = output[0]['payload']
+        for field in ['date', 'message', 'name', 'stack']:
+            assert field in message.keys()
+            
+        assert message['message'] == 'integer division or modulo by zero'
+        assert message['name'] == 'ZeroDivisionError'
+
+    def test_tranform_exception(self):
+        exc_tnfm = StatefulTransform(ExceptionTransform)
+        self.zipline_test_config['transforms'] = [exc_tnfm]
+        
+        zipline = SimulatedTrading.create_test_zipline(
+            **self.zipline_test_config
+        )
+        output, _ = drain_zipline(self, zipline)
+        assert len(output) == 1
+        assert output[0]['prefix'] == 'EXCEPTION'
+        message = output[0]['payload']
+        for field in ['date', 'message', 'name', 'stack']:
+            assert field in message.keys()
+            
+        assert message['message'] == 'An assertion message'
+        assert message['name'] == 'AssertionError'
+
+
     def test_exception_in_init(self):
         # Simulation
         # ----------

zipline/gens/composites.py

@@ -8,14 +8,10 @@ from zipline.gens.sort import date_sort
 from zipline.gens.merge import merge
 from zipline.gens.transform import StatefulTransform
 
-SourceBundle = namedtuple("SourceBundle", ['source', 'args', 'kwargs'])
-TransformBundle = namedtuple("TransformBundle", ['tnfm', 'args', 'kwargs'])
-
 def date_sorted_sources(*sources):
     """
-    Takes an iterable of SortBundles, generating namestrings and
-    initialized datasources for each before piping them into a
-    date_sort.
+    Takes an iterable of sources, generating namestrings and
+    piping their output into date_sort.
     """
 
     for source in sources:
@@ -63,6 +59,7 @@ def merged_transforms(sorted_stream, *transforms):
     # Roundrobin the outputs of our transforms to create a single flat
     # stream.
     to_merge = roundrobin(tnfm_gens, namestrings)
+
     # Pipe the stream into merge.
     merged = merge(to_merge, namestrings)
 

zipline/utils/test_utils.py

@@ -168,3 +168,28 @@ def create_monitor(allocator):
     )
 
     return mon
+
+class ExceptionSource(object):
+    
+    def __init__(self):
+        pass
+
+    def get_hash(self):
+        return "ExceptionSource"
+    
+    def __iter__(self):
+        return self
+
+    def next(self):
+        5 / 0
+        
+class ExceptionTransform(object):
+    
+    def __init__(self):
+        pass
+
+    def get_hash(self):
+        return "ExceptionTransform"
+    
+    def update(self, event):
+        assert False, "An assertion message"