pythonsocketsmockingpython-unittest

How to mock properly recv method in python and use settimeout


I'm trying to use and test sockets in python (I have to use 3.5 version).

Consider this code:

import socket

def send_recv(xml_message):
    try:
        address = ('127.0.0.1', 12000)
        xml_bytes = xml_message.encode(encoding="utf-8")
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client_socket:
            client_socket.settimeout(10) 
            client_socket.connect(address)
            client_socket.sendall(xml_bytes)
            reply = ''
            while True:
                data = client_socket.recv(1024).decode()
                print(data)
                reply += data
                if "<!--MessageEnd-->" in reply or not data:
                    break
            client_socket.shutdown(socket.SHUT_WR)
        return reply
    except ConnectionError as e:
        error_description = "Error connecting to TCP/IP: {}".format(e) 
        print(error_description)
        return None
    except socket.timeout:
        error_description = "Error receiving message. No reply received or message termination not found." 
        print(error_description)
        return None

I tried to test the case where I receive an empty string as a response. I wrote this code:

import unittest
from unittest.mock import patch, MagicMock
from socket_example import only_recv, send_recv

class TestSocketExample(unittest.TestCase):

    @patch('socket_example.socket.socket')
    def test_send_recv(self, mock_socket):
        mock_socket = mock_socket.return_value
        mock_socket.recv.return_value = b''
        #TODO: I have to insert an assert to check the result

I immediately encountered two problems:

  1. If I start the test, the print of the data received

    print(data)

prints the following thing

<MagicMock name='socket().__enter__().recv().decode()' id='137547990704080'>

instead of the empty string.

  1. In the code to be tested I inserted the following code

    client_socket.settimeout(10)

which seems to be ignored because I have to forcefully stop the test.

I don't understand the reason for these code behaviors.


Solution

  • You need to account for the fact that client_socket is the return value of the socket's __enter__ method, not necessarily the actual socket. (In real life it is, but the patching system doesn't know that.)

    mock_socket.__enter__.return_value.recv.return_value = b''
    

    To test the timeout, you need to define a side effect for the fake socket, which makes recv raise an exception instead of returning any value.

    def timeout(*args):
        raise socket.timeout
    
    mock_socket.__enter__.return_value.recv.side_effect = socket.timeout()
    

    The side_effect attribute can take a list of values to return/raise, to simulate one or more successful calls to recv before the timeout finally occurs. You can also assign a function to the attribute, which is called each time the mock is called, allowing you to raise an appropriate exception or return an appropriate value as required for your test.