Fix realtime unittest

mlouielu · Mar 26, 2018 · e1edaf2 · e1edaf2
1 parent 667eca3
commit e1edaf2
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/test/test_realtime.py b/test/test_realtime.py
@@ -7,8 +7,15 @@
 
 class RealtimeTest(unittest.TestCase):
     def test_realtime_field(self):
+        retry = 5
+        for _ in range(retry):
+            stock = realtime.get_raw('2330')
+            if stock['rtcode'] == '0000' and stock['rtmessage'] != 'OK':
+                continue
+            break
+
         self.assertCountEqual(
-            realtime.get_raw('2330').keys(),
+            stock.keys(),
             twstock.mock.get_stock_info('2330').keys())
 
     def test_realtime_get_raw(self):
@@ -24,15 +31,15 @@ def test_realtime_get_blank(self):
         stock = realtime.get('')
 
         self.assertFalse(stock['success'])
-        self.assertIn('rtmessage', stock)
-        self.assertIn('rtcode', stock)
+        self.assertEqual(stock['rtcode'], '0000')
+        self.assertNotEqual(stock['rtmessage'], 'OK')
 
     def test_realtime_get_bad_id(self):
         stock = realtime.get('9999')
 
         self.assertFalse(stock['success'])
-        self.assertIn('rtmessage', stock)
-        self.assertIn('rtcode', stock)
+        self.assertEqual(stock['rtcode'], '5001')
+        self.assertEqual(stock['rtmessage'], 'Empty Query.')
 
         stock = realtime.get(['9999', '8888'])