Make the test more reliable in Python 3.5

shinji-s · Nov 4, 2015 · 3d9e35e · 3d9e35e
1 parent ee0c834
commit 3d9e35e
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tests/unit/heartbeat_tests.py b/tests/unit/heartbeat_tests.py
@@ -152,8 +152,11 @@ def test_new_heartbeat_frame(self):
         self.assertIsInstance(self.obj._new_heartbeat_frame(), frame.Heartbeat)
 
     def test_send_heartbeat_send_frame_called(self):
-        self.obj._send_heartbeat_frame()
-        self.mock_conn._send_frame.assert_called_once()
+        frame_value = self.obj._new_heartbeat_frame()
+        with mock.patch.object(self.obj, '_new_heartbeat_frame') as new_frame:
+            new_frame.return_value = frame_value
+            self.obj._send_heartbeat_frame()
+            self.mock_conn._send_frame.assert_called_once_with(frame_value)
 
     def test_send_heartbeat_counter_incremented(self):
         self.obj._send_heartbeat_frame()