p4-mapreduce

Mocking

This tutorial explains Python mock object library by walking through a project 4 test case. By the end of this tutorial, you should be able to use 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 scripting.

Test case name Description
test_manager_* Tests your Manager with a mock Worker
test_worker_* Tests your Worker with a mock Manager
test_integration_* Tests your Manager and Worker without mocking

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 hard codes 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": 3001,
        "worker_pid": os.getpid(),
    }).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 function we wrote earlier to a mock Manager socket. This means when your Worker uses the socket library to talk to the Manager, it’s really talking to the mock implementation.

Mock Socket

Next, the test connects the message generator function to a mock socket.

def test_worker_01_register(mocker):
    ...
    mockclientsocket = mocker.MagicMock()
    ...
    mockclientsocket.recv.side_effect = manager_message_generator(mock_socket)

Taking a closer look at mockclientsocket.recv.side_effect, we see a sequence of return values for the recv function. The first time your Worker implementation calls mockclientsocket.recv() it will receive the value from the first yield statement in manager_message_generator(). The second time it calls mockclientsocket.recv() it will return the second yield statement, and so on.

When your Worker calls the socket library accept() function, it gets the mock Manager socket instead.

mock_socket = mocker.patch('socket.socket')
mock_socket.return_value.__enter__.return_value.accept.return_value = (
    mockclientsocket,
    ("127.0.0.1", 10000),
)

Run unit under test

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

try:
    mapreduce.worker.Worker(
        manager_port,  # Manager port
        manager_hb_port,  # Manager heartbeat port
        worker_port,  # Worker port
    )
    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', worker_port)),
    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_pid": os.getpid(),
        "worker_port": 3001
    },
]

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.