CODEWITHSUNDEEP
pythonmockingpytestpython-unittest
Philippe Hebert
Philippe Hebert

Reputation: 2028

pytest/unittest: mock.patch function from module?

Given a folder structure like such:

dags/
  **/
    code.py
tests/
  dags/
    **/
      test_code.py
  conftest.py

Where dags serves as the root of the src files, with 'dags/a/b/c.py' imported as 'a.b.c'.

I want to test the following function in code.py:

from dag_common.connections import get_conn
from utils.database import dbtypes

def select_records(
    conn_id: str,
    sql: str,
    bindings,
):
    conn: dbtypes.Connection = get_conn(conn_id)
    with conn.cursor() as cursor:
        cursor.execute(
            sql, bindings
        )
        records = cursor.fetchall()
    return records

But I am faced with the issue that I fail to find a way to patch the get_conn from dag_common.connections. I attempted the following:

(1) Globally in conftest.py

import os
import sys

# adds dags to sys.path for tests/*.py files to be able to import them
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "dags"))

{{fixtures}}

Where I have tested the following replacements for {{fixtures}}:

(1.a) - default

@pytest.fixture(autouse=True, scope="function")
def mock_get_conn():
    with mock.patch("dag_common.connections.get_conn") as mock_getter:
        yield mock_getter

(1.b) - prefixing path with dags

@pytest.fixture(autouse=True, scope="function")
def mock_get_conn():
    with mock.patch("dags.dag_common.connections.get_conn") as mock_getter:
        yield mock_getter

(1.c) - 1.a, with scope="session"

(1.d) - 1.b, with scope="session"

(1.e) - object patching the module itself

@pytest.fixture(autouse=True, scope="function")
def mock_get_conn():
    import dags.dag_common.connections
    mock_getter = mock.MagicMock()
    with mock.patch.object(dags.dag_common.connections, 'get_conn', mock_getter):
        yield mock_getter

(1.f) - 1.a, but using pytest-mock fixture

@pytest.fixture(autouse=True, scope="function")
def mock_get_conn(mocker):
    with mocker.patch("dag_common.connections.get_conn") as mock_getter:
        yield mock_getter

(1.g) - 1.b, but using pytest-mock fixture

(1.h) - 1.a, but using pytest's monkeypatch

@pytest.fixture(autouse=True, scope="function")
def mock_get_conn(mocker, monkeypatch):
    import dags.dag_common.connections
    mock_getter = mocker.MagicMock()
    monkeypatch.setattr(dags.dag_common.connections, 'get_conn', mock_getter)
    yield mock_getter

(2) Locally applying mock.patch in the test/as a decorator

(2.a) - decorator @mock.patch("dag_common.connections.get_conn")

    @mock.patch("dag_common.connections.get_conn")
    def test_executes_sql_with_default_bindings(mock_getter, mock_context):
        # arrange
        sql = "SELECT * FROM table"
        records = [RealDictRow(col1=1), RealDictRow(col1=2)]
        mock_conn = mock_getter.return_value
        mock_cursor = mock_conn.cursor.return_value
        mock_cursor.execute.return_value = records
        # act
        select_records(conn_id="orca", sql=sql, ) # ...
        # assert
        mock_cursor.execute.assert_called_once_with(
            sql, # ...
        )

(2.b) - (2.a) but with "dags." prefix

(2.c) - context manager

    def test_executes_sql_with_default_bindings(mock_context):
        # arrange
        sql = "SELECT * FROM table"
        records = [RealDictRow(col1=1), RealDictRow(col1=2)]
        with mock.patch("dag_common.connections.get_conn") as mock_getter:
            mock_conn = mock_getter.return_value
            mock_cursor = mock_conn.cursor.return_value
            mock_cursor.execute.return_value = records
            # act
            select_records(conn_id="orca", sql=sql, ) # ...
            # assert
            mock_cursor.execute.assert_called_once_with(
                sql, # ...
            )

(2.d) - (2.c) but with "dags." prefix

Conclusion

But alas, no matter what solution I pick, the function-to-be-mocked still gets called. I made sure to attempt each solution separatedly from each other, and to kill/clear/restart my pytest-watch process in between attempts.

I feel like this may be related to me meddling with sys.path in conftest.py, because outside of this I feel like I have exhausted all possibilities.

Any idea how I can solve this?

Upvotes: 16

Views: 32013

Answers (2)

Giorgos Myrianthous
Giorgos Myrianthous

Reputation: 39830

Since the example you are showcasing is relevant to Apache Airflow, another way you can actually use in order to mock Airflow Connections and/or Variables is by mocking the corresponding environment variables.

For example,

from unittest import mock

import pytest
from airflow.models.connection import Connection


@pytest.fixture(autouse=True, scope='function')
def mock_my_conn():
    conn = Connection(conn_type='', login='', host='')
    yield conn.get_uri()


def test_my_dag(mock_my_conn):
    with mock.patch.dict('os.environ', AIRFLOW_CONN_MY_CONN=mock_my_conn):
        ...

Or if you want to mock multiple connections in a single test case:

def test_my_dag(mock_my_conn):
    with mock.patch.dict(
        'os.environ', {
            'AIRFLOW_CONN_MY_CONN': mock_my_conn,
            'AIRFLOW_CONN_ANOTHER_CONN': mock_my_conn,
        }
    ):
        ...

Upvotes: 1

Jarek Potiuk
Jarek Potiuk

Reputation: 20077

Yeah. I also fought with this initially when I learned patching and mocking and know how frustrating it is as you seem to be doing everything right, but it does not work. I sympathise with you!

This is actually how mocking of imported stuff works, and once you realise it, it actually makes sense.

The problem is that import works in the way that it makes the imported module available in the context of where your import is.

Lets' assume your code.py module is in 'my_package' folder. Your code is available then as my_package.code. And once you use from dag_common.connections import get_conn in code module - the imported get_conn becomes available as .... my_package.code.get_conn

And in this case you need to patch my_package.code.get_conn not the original package you imported get_conn from.

Once you realise this, patching becomes much easier.

Upvotes: 38

Related Questions