fix contact tcp tests

mitre · Dec 9, 2023 · 60412a7 · 60412a7
1 parent 599419d
commit 60412a7
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/tests/contacts/test_contact_tcp.py b/tests/contacts/test_contact_tcp.py
@@ -9,28 +9,34 @@
 
 class TestTcpSessionHandler:
 
-    def test_refresh_with_socket_errors(self, event_loop):
+    def test_refresh_with_socket_errors(self, event_loop, async_return):
         handler = TcpSessionHandler(services=None, log=logger)
 
         session_with_socket_error = mock.Mock()
         session_with_socket_error.connection.send.side_effect = socket.error()
 
+        standard_session = mock.Mock()
+        standard_session.connection.send.return_value = async_return(True)
+
         handler.sessions = [
             session_with_socket_error,
             session_with_socket_error,
-            mock.Mock()
+            standard_session
         ]
 
         event_loop.run_until_complete(handler.refresh())
         assert len(handler.sessions) == 1
         assert all(x is not session_with_socket_error for x in handler.sessions)
 
-    def test_refresh_without_socket_errors(self, event_loop):
+    def test_refresh_without_socket_errors(self, event_loop, async_return):
+        standard_session = mock.Mock()
+        standard_session.connection.send.return_value = async_return(True)
+
         handler = TcpSessionHandler(services=None, log=logger)
         handler.sessions = [
-            mock.Mock(),
-            mock.Mock(),
-            mock.Mock()
+            standard_session,
+            standard_session,
+            standard_session
         ]
 
         event_loop.run_until_complete(handler.refresh())