Example of Mocking a Socket

Server Code:

​

import socket

​

def echo_server(host, port):

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:

        server_socket.bind((host, port))

        server_socket.listen()

        conn, addr = server_socket.accept()

        with conn:

            print('Connected by', addr)

            while True:

                data = conn.recv(1024)

                if not data:

                    break

                conn.sendall(data)

Mocking the Server for Testing:

import unittest

from unittest.mock import patch, MagicMock

​

class TestEchoServer(unittest.TestCase):

    @patch('socket.socket')

    def test_echo_server(self, mock_socket):

        # Create a mock socket instance

        mock_conn = MagicMock()

        mock_socket_instance = MagicMock()

        mock_socket_instance.accept.return_value = (mock_conn, ('127.0.0.1', 12345))

        mock_socket.return_value = mock_socket_instance

​

        # Simulate the data to be received

        mock_conn.recv.side_effect = [b'Hello', b'', b'']

        # Call the function to be tested

        echo_server('localhost', 12345)

        # Assertions to check if the data was sent back correctly

        mock_conn.sendall.assert_called_with(b'Hello')

​

if __name__ == '__main__':

    unittest.main()

Mocking a Client:

def echo_client(host, port, message):

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client_socket:

        client_socket.connect((host, port))

        client_socket.sendall(message)

        data = client_socket.recv(1024)

        return data

​

class TestEchoClient(unittest.TestCase):

    @patch('socket.socket')

    def test_echo_client(self, mock_socket):

        mock_socket_instance = MagicMock()

        mock_socket.return_value = mock_socket_instance

        mock_socket_instance.recv.return_value = b'Echo: Hello'

​

        response = echo_client('localhost', 12345, b'Hello')

        mock_socket_instance.connect.assert_called_with(('localhost', 12345))

        mock_socket_instance.sendall.assert_called_with(b'Hello')

        self.assertEqual(response, b'Echo: Hello')

​

if __name__ == '__main__':

    unittest.main()