Merge pull request #806 from tlsfuzzer/new-mock

fix called_once_with() asserts
tlsfuzzer · Feb 10, 2023 · 7a7011a · 7a7011a
2 parents b01bfdb + 31f3463
commit 7a7011a
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 12 deletions.
diff --git a/tests/test_tlsfuzzer_expect.py b/tests/test_tlsfuzzer_expect.py
@@ -3988,7 +3988,7 @@ def test_process_with_matching_type(self):
 
         exp.process(state, ku)
 
-        state.msg_sock.calcTLS1_3PendingState.called_once_with(
+        state.msg_sock.calcTLS1_3KeyUpdate_sender.assert_called_once_with(
             cipher, cats, sats)
         self.assertIs(state.key['server application traffic secret'], ret)
 

diff --git a/tests/test_tlsfuzzer_messages.py b/tests/test_tlsfuzzer_messages.py
@@ -116,7 +116,7 @@ def test_process(self):
         close = Close()
         close.process(state)
 
-        state.msg_sock.sock.close.called_once_with()
+        state.msg_sock.sock.close.assert_called_once_with()
 
 class TestTCPBufferingEnable(unittest.TestCase):
     def test___init__(self):
@@ -458,9 +458,8 @@ def test_process(self):
 
         msg_gen.process(state)
 
-        self.assertTrue(
-            state.msg_sock._recordSocket.
-            sock.send.called_once_with(b'some data'))
+        state.msg_sock._recordSocket.sock.send.assert_called_once_with(
+            b'some data')
 
 
 class TestPlaintextMessageGenerator(unittest.TestCase):

diff --git a/tests/test_tlsfuzzer_runner.py b/tests/test_tlsfuzzer_runner.py
@@ -190,10 +190,11 @@ def test_run_with_expect_node(self):
 
         runner.run()
 
-        internal_message = messages.Message(msg[0].type, msg[1].bytes)
-
-        node.is_match.called_once_with(internal_message)
-        node.process.called_once_with(runner.state, internal_message)
+        # as the message they're called with is generated inside the runner
+        # it will be a different object every time, so just assert that
+        # the methods were called
+        node.is_match.assert_called_once()
+        node.process.assert_called_once()
 
     def test_run_with_expect_and_closed_socket(self):
         node = ExpectClose()
@@ -279,7 +280,7 @@ def test_run_with_expect_node_and_unexpected_message(self):
         with self.assertRaises(AssertionError):
             runner.run()
 
-        runner.state.msg_sock.sock.close.called_once_with()
+        runner.state.msg_sock.sock.close.assert_called_once_with()
 
     def test_run_with_generate_and_unexpected_closed_socket(self):
         node = mock.MagicMock()

diff --git a/tlsfuzzer/runner.py b/tlsfuzzer/runner.py
@@ -229,8 +229,6 @@ def run(self):
                     node = next((proc for proc in node.get_all_siblings()
                                  if proc.is_match(msg)), None)
                     if node is None:
-                        # since we're aborting, the user can't clean up
-                        self.state.msg_sock.sock.close()
                         raise AssertionError("Unexpected message from peer: " +
                                              guess_response(\
                                                  msg.contentType,