p4-mapreduce

Mocking

This tutorial explains the Python unittest.mock module by walking through a project 4 test case. By the end of this tutorial, you should be able to step through a test case that uses mocking and understand how it provides input to your code.

A mock object replaces part of your program for testing. For example, your MapReduce framework has a Manager and a Worker that communicate over the network. With mocking, we can independently test a Worker implementation by mocking the responses from the Manager.

Overview

There are 3 types of test cases in project 4 apart from code style and shell scripting.

Example

Let’s take a look at test_worker_01.py, which verifies that your Worker correctly registers with the Manager. This test uses your Worker code and mocks the messages of the Manager.

Message Generator

At the beginning of test_worker_01.py, we see a function that hardcodes the messages sent by the Manager in this test. Each yield statement transfers control back to your Worker code.

def manager_message_generator(mock_socket):
    """Fake Manager messages."""
    # First message
    utils.wait_for_register_messages(mock_socket)
    yield json.dumps({
        "message_type": "register_ack",
        "worker_host": "localhost",
        "worker_port": 6001,
    }).encode('utf-8')
    yield None

    # Shutdown
    yield json.dumps({
        "message_type": "shutdown",
    }).encode('utf-8')
    yield None

Later, we’ll see that the test connects manager_message_generator() to a mock client socket that would normally receive messages from the Manager. Whenever your Worker calls a function on a socket, it’s calling it on a mock socket instead of a real one.

Mock Socket

The test connects the message generator function to a mock socket.

We first mock the socket library’s socket class. We can then mock anything produced by the socket class, like the sendall() function.

def test_register(mocker):
    # Mock the socket library socket class
    mock_socket = mocker.patch("socket.socket")

    # sendall() records messages
    mock_sendall = mock_socket.return_value.__enter__.return_value.sendall

Mock the socket accept() function to return a mock client socket instead of a of real one.

    # accept() returns a mock client socket
    mock_clientsocket = mocker.MagicMock()
    mock_accept = mock_socket.return_value.__enter__.return_value.accept
    mock_accept.return_value = (mock_clientsocket, ("127.0.0.1", 10000))

Finally, connect our manager_message_generator() function to the mock socket. When your code calls recv() it will get the next value yielded by the generator.

    # recv() returns values generated by manager_message_generator()
    mock_recv = mock_clientsocket.recv
    mock_recv.side_effect = manager_message_generator(mock_sendall)

Run unit under test

The test then runs your Worker implementation. Each time the Worker calls recv(), manager_message_generator() yields another message.

try:
    mapreduce.worker.Worker(
        host="localhost",
        port=6001,
        manager_host="localhost",
        manager_port=6000,
    )
    utils.wait_for_threads()
except SystemExit as error:
    assert error.code == 0

When the Worker shuts down, the test case continues.

Verify function calls

A mock object keeps track of function calls. At the end of the test, we verify that your Worker called the correct socket functions. This code checks that the Worker called 4 methods: the constructor, setsockopt, bind, and listen with correct arguments.

mock_socket.assert_has_calls([
    mocker.call(socket.AF_INET, socket.SOCK_STREAM),
    mocker.call().__enter__().setsockopt(
        socket.SOL_SOCKET,
        socket.SO_REUSEADDR,
        1,
    ),
    mocker.call().__enter__().bind(('localhost', 6001)),
    mocker.call().__enter__().listen(),
], any_order=True)

Verify messages

A mock object keeps track of function call arguments. The test gets a list of the messages your Worker sent to the mock Manager. Then, it verifies that the messages look OK.

messages = utils.get_messages(mock_socket)
messages = utils.filter_not_heartbeat_messages(messages)
assert messages == [
    {
        "message_type": "register",
        "worker_host": "localhost",
        "worker_port": 6001
    },
]

The utils.get_messages() function collects the messages sent from your Worker to the mock Manager using the socket library sendall().

Further reading

Check out this mocking tutorial if you would like to learn more.

Acknowledgments

Original document written by Zachary Papanastasopoulos and Andrew DeOrio awdeorio@umich.edu.

This document is licensed under a Creative Commons Attribution-NonCommercial 4.0 License. You’re free to copy and share this document, but not to sell it. You may not share source code provided with this document.