Python - Mocking nested DB calls

Question

I'm trying to write unit test for some_func function in some_func.py below. I do not want to connect to any database during this tests and I want to mock out any calls to DB.

Since the actual db calls are somewhat nested, i'm not able to get this to work. How do i patch any db interaction in this case?

db_module.py

import MySQLdb

from random_module.config.config_loader import config
from random_module.passwd_util import get_creds

class MySQLConn:

    _conn = None

    def __init__(self):
        self._conn = self._get_rds_connection()

    def get_conn(self):
        return self._conn

    def _get_rds_connection(self):
        """
        Returns a conn object after establishing connection with the MySQL database

        :return: obj: conn
        """

        try:

            logger.info("Establishing connection with MySQL")

            username, password = get_creds(config['mysql_dev_creds'])

            connection = MySQLdb.connect(
                host=config['host'],
                user=username,
                passwd=password,
                db=config['db_name'],
                port=int(config['db_port']))

            connection.autocommit = True

        except Exception as err:
            logger.error("Unable to establish connection to MySQL")
            raise ConnectionAbortedError(err)

        if (connection):
            logger.info("Connection to MySQL successful")
            return connection
        else:
            return None

db_mysql = MySQLConn()

some_func.py

from random_module.utils.db_module import db_mysql

def some_func():

    try: 

        db_conn = db_mysql.get_conn()
        db_cursor = db_conn.cursor()

        results = db_cursor.execute("SELECT * FROM some_table")
        results = db_cursor.fetchall()

        result_set = []

        for row in results:

            result_set.insert(i, row['x'])
            result_set.insert(i, row['y'])

    except Exception:
        logger.error("some error")

    return result_set

dir struct -

src
├──pkg
|   ├── common
|      |___ some_func.py
|      |___ __init__.py
|       
|   ├── utils
|      |___ db_module.py
|      |___ __init__.py
|
| __init__.py

Dan · Accepted Answer · 2018-04-28 23:56:11Z

You need to mock this line: db_conn = db_mysql.get_conn()

The return value of the get_conn method is what you're interested in.

from random_module.utils.db_module import db_mysql

@mock.patch.object(db_mysql, 'get_conn')
def test_some_func(self, mock_get):
    mock_conn = mock.MagicMock()
    mock_get.return_value = mock_conn
    mock_cursor = mock.MagicMock()
    mock_conn.cursor.return_value = mock_cursor

    expect = ...
    result = some_func()

    self.assertEqual(expect, result)
    self.assertTrue(mock_cursor.execute.called)

As you can see there is a lot of complexity setting up these mocks. That is because you are instantiating objects inside of your function. A better approach would be to refactor your code to inject the cursor because the cursor is the only thing relevant to this function. An even better approach would be to create a database fixture to test the function actually interacts with the database correctly.

Will Keeling · Accepted Answer · 2018-04-23 09:28:40Z

0

You should mock db_mysql in some_module.py and then assert that the expected calls were made after some_func() executes.

from unittest TestCase
from unittest.mock import patch
from some_module import some_func

class SomeFuncTest(TestCase):

    @patch('some_module.db_mysql')
    def test_some_func(self, mock_db):
        result_set = some_func()

        mock_db.get_conn.assert_called_once_with()
        mock_cursor = mock_db.get_conn.return_value
        mock_cursor.assert_called_once_with()

        mock_cursor.execute.assert_called_once_with("SELECT * FROM some_table")
        mock_cursor.fetchall.return_value = [
            {'x': 'foo1', 'y': 'bar1'},
            {'x': 'foo2', 'y': 'bar2'}
        ]
        mock_cursor.fetchall.assert_called_once_with()

        self.assertEqual(result_set, ['foo1', 'bar1', 'foo2', 'bar2'])

answered Apr 23, 2018 at 9:28

Will Keeling

23.2k4 gold badges58 silver badges64 bronze badges

8 Comments

Naveen Over a year ago

I edited the question as i have the some_func.py module and some_func method same. When i try to patch i'm doing @mock.patch('pkg.common.some_func.db_mysql'). But this is resulting in AttributeError: 'module' object has no attribute 'some_func'. I even have the init.py listed in the common directory.

Will Keeling Over a year ago

The package module is called __init__.py. You have one of these in pkg and common?

Naveen Over a year ago

Updated the dir structure. I have a __init__.py in every folder.

Will Keeling Over a year ago

Do you have a common.py alongside your common package? It seems the patch is getting confused there somehow.

Naveen Over a year ago

No I don't. Since db_mysql is getting instantiated in a different class, do you think that could be the culprit?

|

Collectives™ on Stack Overflow

Python - Mocking nested DB calls

2 Answers 2

Comments

8 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

8 Comments

Your Answer

Sign up or log in

Post as a guest

Related