Python mocking function from imported module

Question

How to write unit test case for following code. The code just returns a db connection. I can't create a real db connection. So, i am dependent on object mocking.

File : code.py

import snowflake.connector
import logging
import simplejson as json


def get_database_connection(user, password, account, warehouse, database=None, schema=None):
    logging.info("DB Connection : START : " + __name__)
    try:
        return snowflake.connector.connect(
            user=user,
            password=password,
            account=account,
            warehouse=warehouse,
            database=database,
            schema=schema
        )
    except Exception:
        logging.exception("Exception while creating connection")

test_code.py

class SnowflakeDbTest(unittest.TestCase):

    def test_get_database_connection(self):
         snowflake = Mock()
         snowflake.connector.return_value = Mock()
         snowflake.connector.connect.return_value = "dbconnection"
         self.assertEqual("dbconnection", code.get_database_connection("user", "password", "account", "warehouse", database=None, schema=None))

Solution

    @patch('code.snowflake.connector')
    def test_get_database_connection(self, mock_connection):
         mock_connection.connect.return_value = "dbconnection"
         self.assertEqual("dbconnection", snowflakedb.get_database_connection("user", "password", "account", "warehouse", database=None, schema=None))

How to mock 'snowflake.connector' which is getting imported in 'code.py'

Answer 1

patching is the solution

    @patch('code.snowflake.connector')
    def test_get_database_connection(self, mock_connection):
         mock_connection.connect.return_value = "dbconnection"
         self.assertEqual("dbconnection", snowflakedb.get_database_connection("user", "password", "account", "warehouse", database=None, schema=None))