From e22a66b1ae6a1f90790d70acfd20c29c90e37d0f Mon Sep 17 00:00:00 2001
From: fawce <fawce@quantopian.com>
Date: Sun, 12 Feb 2012 08:02:56 -0500
Subject: [PATCH] moved unit test asserts into the test client

---
 qsim/test/client.py         | 68 ++++++++++++++++++++++---------------
 qsim/test/test_messaging.py |  9 ++---
 2 files changed, 42 insertions(+), 35 deletions(-)

diff --git a/qsim/test/client.py b/qsim/test/client.py
index a23b3f1f..d9350d70 100644
--- a/qsim/test/client.py
+++ b/qsim/test/client.py
@@ -9,41 +9,53 @@ import qsim.messaging as qmsg
 
 class TestClient(object):
     
-    def __init__(self):
-        self.address        = None
-        self.sync           = None
-        self.received_count = 0
+    def __init__(self, utest, expected_msg_count=0):
+        self.address            = None
+        self.sync               = None
+        self.received_count     = 0
+        self.expected_msg_count = expected_msg_count
+        self.ERROR              = False
+        self.utest              = utest
         
     def run(self):
         
-        qutil.LOGGER.info("running the client")
-        self.context = zmq.Context()
+        try:
+            qutil.LOGGER.info("running the client")
+            self.context = zmq.Context()
         
-        self.data_feed = self.context.socket(zmq.PULL)
+            self.data_feed = self.context.socket(zmq.PULL)
 
-        qutil.LOGGER.info("connecting to {address}".format(address=self.address))
-        self.data_feed.connect(self.address)
+            qutil.LOGGER.info("connecting to {address}".format(address=self.address))
+            self.data_feed.connect(self.address)
         
-        self.sync.confirm()
+            self.sync.confirm()
         
-        qutil.LOGGER.info("Starting the client loop")
+            qutil.LOGGER.info("Starting the client loop")
 
-        prev_dt = None
-        while True:
-            msg = self.data_feed.recv()
-            if(msg == "DONE"):
-                qutil.LOGGER.info("DONE!")
-                break
-            self.received_count += 1
-            event = json.loads(msg)
-            if(prev_dt != None):
-                if(not event['dt'] >= prev_dt):
-                    raise Exception("message arrived out of order: {date} after {prev}".format(date=event['dt'], prev=prev_dt))
+            prev_dt = None
+            while True:
+                msg = self.data_feed.recv()
+                if(msg == "DONE"):
+                    qutil.LOGGER.info("DONE!")
+                    break
+                self.received_count += 1
+                event = json.loads(msg)
+                if(prev_dt != None):
+                    if(not event['dt'] >= prev_dt):
+                        raise Exception("message arrived out of order: {date} after {prev}".format(date=event['dt'], prev=prev_dt))
             
-            prev_dt = event['dt']
-            if(self.received_count % 100 == 0):
-                qutil.LOGGER.info("received {n} messages".format(n=self.received_count))
+                prev_dt = event['dt']
+                if(self.received_count % 100 == 0):
+                    qutil.LOGGER.info("received {n} messages".format(n=self.received_count))
             
-        qutil.LOGGER.info("received {n} messages".format(n=self.received_count))
-        self.data_feed.close()
-        self.context.term()
\ No newline at end of file
+            qutil.LOGGER.info("received {n} messages".format(n=self.received_count))
+        except:
+            self.ERROR = True
+            qutil.LOGGER.exception("Error in test client.")
+        finally:
+            self.data_feed.close()
+            self.context.term()
+        
+        self.utest.assertEqual(self.expected_msg_count, self.received_count, 
+                        "The client should have received ({n}) the same number of messages as the feed sent ({m})."
+                            .format(n=self.received_count, m=self.expected_msg_count))
\ No newline at end of file
diff --git a/qsim/test/test_messaging.py b/qsim/test/test_messaging.py
index 3653d983..6a74e164 100644
--- a/qsim/test/test_messaging.py
+++ b/qsim/test/test_messaging.py
@@ -28,16 +28,13 @@ class MessagingTestCase(unittest.TestCase):
         ret1 = RandomEquityTrades(133, "ret1", 400)
         ret2 = RandomEquityTrades(134, "ret2", 400)
         sources = {"ret1":ret1, "ret2":ret2}
-        client = TestClient()
+        client = TestClient(self, expected_msg_count=800)
         sim = Simulator(sources, {}, client)
         sim.launch()
               
         self.assertEqual(sim.feed.data_buffer.pending_messages(), 0, 
                         "The feed should be drained of all messages, found {n} remaining."
                             .format(n=sim.feed.data_buffer.pending_messages()))
-        self.assertEqual(800, client.received_count, 
-                        "The client should have received ({n}) the same number of messages as the feed sent ({m})."
-                            .format(n=client.received_count, m=800))
     
     
     def test_merged_to_client(self):
@@ -52,12 +49,10 @@ class MessagingTestCase(unittest.TestCase):
         mavg1 = MovingAverage("mavg1", 30)
         mavg2 = MovingAverage("mavg2", 60)
         transforms = {"mavg1":mavg1, "mavg2":mavg2}
-        client = TestClient()
+        client = TestClient(self, expected_msg_count=800)
         sim = Simulator(sources, {}, client)
         sim.launch()
         
         
         self.assertEqual(sim.feed.data_buffer.pending_messages(), 0, "The feed should be drained of all messages.")
-        self.assertEqual(800, client.received_count, 
-                        "The client should have received the same number of messages as the feed sent.")