From c64175e2e59cf2cc440754f71fefef9f987183cd Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Wed, 8 Aug 2012 18:44:45 -0400
Subject: [PATCH] added tests for exception messages in datasources and
 transforms

---
 tests/test_exception_handling.py | 38 +++++++++++++++++++++++++++++++-
 zipline/gens/composites.py       |  9 +++-----
 zipline/utils/test_utils.py      | 25 +++++++++++++++++++++
 3 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/tests/test_exception_handling.py b/tests/test_exception_handling.py
index f547ecc9..f604466a 100644
--- a/tests/test_exception_handling.py
+++ b/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
         # ----------
diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py
index 7fc5d71b..b3fa7576 100644
--- a/zipline/gens/composites.py
+++ b/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)
 
diff --git a/zipline/utils/test_utils.py b/zipline/utils/test_utils.py
index 02ac4c69..c31b2f25 100644
--- a/zipline/utils/test_utils.py
+++ b/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"