From e048e8bc352631747dcc757b70631708d72a686b Mon Sep 17 00:00:00 2001
From: scottsanderson <ssanderson@quantopian.com>
Date: Sat, 28 Jul 2012 19:04:50 -0400
Subject: [PATCH] added done message to SpecificEquity

---
 zipline/gens/feed.py      | 3 ++-
 zipline/gens/test_feed.py | 4 +---
 zipline/gens/tradegens.py | 6 ++++--
 zipline/gens/utils.py     | 3 +++
 4 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/zipline/gens/feed.py b/zipline/gens/feed.py
index 93e03371..0528d7ee 100644
--- a/zipline/gens/feed.py
+++ b/zipline/gens/feed.py
@@ -52,7 +52,8 @@ def FeedGen(stream_in, source_ids):
             message = pop_oldest(sources)
             assert_feed_protocol(message)
             yield message
-
+    
+    import nose.tools; nose.tools.set_trace()
     # We should have only a done message left in each queue.    
     for queue in sources.itervalues():
         assert len(queue) == 1, "Bad queue in FeedGen on exit: %s" % queue
diff --git a/zipline/gens/test_feed.py b/zipline/gens/test_feed.py
index 9b77f9c8..12cf4bc5 100644
--- a/zipline/gens/test_feed.py
+++ b/zipline/gens/test_feed.py
@@ -184,10 +184,8 @@ class FeedGenTestCase(TestCase):
         sources = (source_a, source_b, source_c)
         source_ids = [id_a, id_b, id_c]
         
-        import nose.tools; nose.tools.set_trace()
         feed_out = PreTransformLayer(sources, source_ids)
-        for i in feed_out:
-            print i
+        l = list(feed_out)
     
 def mock_data_unframe(source_id, dt, type):
     event = ndict()
diff --git a/zipline/gens/tradegens.py b/zipline/gens/tradegens.py
index 8327aa63..35100f8e 100644
--- a/zipline/gens/tradegens.py
+++ b/zipline/gens/tradegens.py
@@ -3,7 +3,7 @@ from itertools import chain, repeat, cycle, ifilter, izip
 from datetime import datetime, timedelta
 
 from zipline.utils.factory import create_trade
-from zipline.gens.utils import hash_args
+from zipline.gens.utils import hash_args, mock_done
 
 def date_gen(start = datetime(2012, 6, 6, 0),
              delta = timedelta(minutes = 1), 
@@ -77,7 +77,9 @@ def SpecificEquityTrades(count = 500,
     else:
         filtered = unfiltered
 
-    return filtered
+    # Add a done message to the end of the stream.
+    out = chain(filtered, iter([mock_done(namestring)]))
+    return out    
 
 def RandomEquityTrades(count = 500, sids = [1,2], filter = None):
     dates = fuzzy_dates(500)
diff --git a/zipline/gens/utils.py b/zipline/gens/utils.py
index 8966c638..8ed6ea9d 100644
--- a/zipline/gens/utils.py
+++ b/zipline/gens/utils.py
@@ -16,6 +16,9 @@ def mock_raw_event(sid, dt):
     }
     return event
 
+def mock_done(source_id):
+    return ndict({'dt': "DONE", "source_id" : source_id, 'type' : 0})
+              
 def alternate(g1, g2):
     """Specialized version of roundrobin for just 2 generators."""
     for e1, e2 in izip_longest(g1, g2):