From 9deaefe6e60afe11642b77785f1f825f02b3f5c2 Mon Sep 17 00:00:00 2001
From: fawce <fawce@quantopian.com>
Date: Mon, 30 Jul 2012 13:52:38 -0400
Subject: [PATCH] fixes for tests

---
 tests/test_exception_handling.py |  2 +-
 zipline/core/component.py        |  4 ++--
 zipline/utils/test_utils.py      | 34 +++++++++++++++++++++-----------
 3 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/tests/test_exception_handling.py b/tests/test_exception_handling.py
index 013a9624..6b67c9db 100644
--- a/tests/test_exception_handling.py
+++ b/tests/test_exception_handling.py
@@ -3,7 +3,7 @@ import zmq
 from unittest2 import TestCase
 from collections import defaultdict
 
-from zipline.test_algorithms import ExceptionAlgorithm, NoopAlgorithm
+from zipline.test_algorithms import ExceptionAlgorithm
 from zipline.finance.trading import SIMULATION_STYLE
 from zipline.core.devsimulator import AddressAllocator
 from zipline.lines import SimulatedTrading
diff --git a/zipline/core/component.py b/zipline/core/component.py
index 8a399cb2..0e1d04f8 100644
--- a/zipline/core/component.py
+++ b/zipline/core/component.py
@@ -520,9 +520,9 @@ class Component(object):
                 # notice arrives, and we can assume other zipline
                 # components have broken out of their message
                 # loops.
-                for i in xrange(100):
+                for i in xrange(PARAMETERS.MAX_COMPONENT_WAIT):
                     self.heartbeat(timeout=1000)
-                log.warn("{id} Never heard back from monitor."\
+                log.warn("{id} never heard back from monitor."\
                         .format(id=self.get_id))
             except:
                 log.exception("Exception waiting for controller reply")
diff --git a/zipline/utils/test_utils.py b/zipline/utils/test_utils.py
index 8d0d93c2..6655c265 100644
--- a/zipline/utils/test_utils.py
+++ b/zipline/utils/test_utils.py
@@ -65,18 +65,36 @@ def drain_zipline(test, zipline):
     assert isinstance(test.zipline_test_config, dict)
     assert test.zipline_test_config['results_socket'], \
             "need to specify a socket address for logs/perf/risk"
-    test.receiver = test.ctx.socket(zmq.PULL)
-    test.receiver.bind(test.zipline_test_config['results_socket'])
+    test.receiver = create_receiver(
+        test.zipline_test_config['results_socket'],
+        test.ctx
+    )
     # Bind and connect are asynch, so allow time for bind before
     # starting the zipline (TSC connects internally).
     time.sleep(1)
+
     # start the simulation
     zipline.simulate(blocking=False)
+    output, transaction_count = drain_receiver(test.receiver)
+    # some processes will exit after the message stream is
+    # finished. We block here to avoid collisions with subsequent
+    # ziplines.
+    for process in zipline.sim.subprocesses:
+        process.join()
 
+    return output, transaction_count
+
+def create_receiver(socket_addr, ctx):
+    receiver = ctx.socket(zmq.PULL)
+    receiver.bind(socket_addr)
+
+    return receiver
+
+def drain_receiver(receiver):
     output = []
     transaction_count  = 0
     while True:
-        msg = test.receiver.recv()
+        msg = receiver.recv()
         if msg == str(zp.CONTROL_PROTOCOL.DONE):
             break
         else:
@@ -88,14 +106,8 @@ def drain_zipline(test, zipline):
             elif update['prefix'] == 'EXCEPTION':
                 break
 
-    test.receiver.close()
-    del test.receiver
-
-    # some processes will exit after the message stream is
-    # finished. We block here to avoid collisions with subsequent
-    # ziplines.
-    for process in zipline.sim.subprocesses:
-        process.join()
+    receiver.close()
+    del receiver
 
     return output, transaction_count