What I Learned at Work this Week: MagicMock

Photo by Egor Kamelev: https://www.pexels.com/photo/close-up-photography-of-snowflake-813872/

The Python testing continues! This past week, one of my fellow Solutions Engineers shared a PR that mocks a DB, which is new for a lot of us. It’ll be important to write more of those tests in the future, so I’ve got to do the legwork to try and understand everything in there. One of the first lines of code uses MagicMock, so let’s learn about what that means.

unittest.mock.MagicMock

In my last post, I wrote about Mock, which we can use to replace objects when writing tests in Python. MagicMock is a sublcass of Mock, so it does all the same things, but includes some default results for common methods:

credit docs.python.org

In the documentation, they show the benefit of these methods with an example:

>>> mock = MagicMock()
>>> int(mock)
1
>>> len(mock)
0
>>> list(mock)
[]
>>> object() in mock
False

MagicMock has default results when passed as an argument in methods like int or len. Depending on what we’re testing, we can replace these values so that our mocked object can act as a needed patch. The documentation uses __eq__ as the example here:

>>> MagicMock() == 3
False
>>> MagicMock() != 3
True
>>> mock = MagicMock()
>>> mock.__eq__.return_value = True
>>> mock == 3
True

Note that we don’t simply set the value of our object, but instead set the return_value property of the associated method. When I looked back at the PR, I was curious whether this functionality was the reason my teammate had used MagicMock. I didn’t see anything that looked like this code, so I looked up whether there was a preference for Mock vs MagicMock as a default.

Stack Overflow pointed me to an answer from the author of Mock, Michael Foord. Foord points out that the default values of MagicMock could get in the way of testing if you’re trying to mock an object where a certain property is undefined. It could also cause an issue if we inadvertently used it to set the wrong default value. With that said, the official Python documentation does literally say as the MagicMock is the more capable class it makes a sensible one to use by default. I was originally going to comment on the PR and suggest changing to Mock, but it seems that there’s arguments for both sides about whether it’s better to use a capable object with a lot of built-in values or more of a blank slate.

Mocking a Snowflake Connection

My original intent for this post was not to examine the philosophical questions behind the MagicMock object, but to understand how it was used to mock a call to a database. Let’s check out some code:

@patch("snowflake.connector.connect")
@mock.create_autospec
def test_snowflake_query(mocker):
    mock_conn = MagicMock()
    mocker.patch("snowflake.connector.connect", return_value=mock_conn)

    mock_cur = MagicMock()
    mock_conn.cursor.return_value = mock_cur

    mock_fetchall = MagicMock(return_value=expected_return_value)
    mock_cur.fetchall = mock_fetchall

Here’s the first part of the test, which changes made so as not to share exact private code publicly. snowflake.connector.connect comes from the Snowflake Connector library, and is a method that will allow us to use Python to query Snowflake. My teammate wrote a patch that mocks the method, so there’s no need to worry about hitting a real DB. When we test module that uses Snowflake connection, it’ll reference the MagicMock instead.

Our Snowflake interaction involves three distinct steps. We’ve covered the first by mocking the connection, so next we move on to mocking the cursor itself. In Python, the cursor object is used to execute SQL statements. If we’re not using a real DB connection, we can’t use a real cursor either, so that gets mocked as well.

The third step is fetching the data returned from the query. Mock DB, mock query, mock fetch! We can give the mock fetch an expected return value and then confirm that the fetch was run. In the rest of the code, my teammate wrote tests to check:

1. that running the module will trigger an attempt to connect to Snowflake (snowflake.connector.connect.assert_called_once_with(… )
2. the cursor is used (mock_cur.execute.assert_called_once_with(… )
3. fetchall is called and returns the expected value:

mock_fetchall.assert_called_once()

assert result == expected_return_value

Testing what we can

It’s impractical to make a real SQL query or API call when running a test. Though it might be the most clear way to confirm something is working, it can be costly, time-consuming, and inconsistent. Instead, we should do what my teammate did and seek to identify what we can control and build assurances that those parts of the code will work. Will a call be made to the DB? What happens after data is queried? If the fetch fails, how does our code react? What if the data is malformed? These are all things we can mock and test, and thanks to this new code, it’ll be easier for all of my coworkers to develop those tests.

Sources

• The Mock Class, python docs
• MagicMock, python docs
• Mock vs MagicMock, Stack Overflow
• Testing with Mock, Michael Foord
• Mock Patching Methods, python docs