Mocking Raise Exception in function with Pytest

Question

I have the following function get_postgres_connection. I'm trying to run a unittest with test_get_postgres_connection_unsuccess to hit the exception.

def get_postgres_connection():
        
        try:
            conn = psycopg2.connect(dbname = POSTGRES_DATABASE,
                                user = POSTGRES_USER,
                                password = POSGRES_PASSWORD,
                                host = POSTGRES_HOST,
                                port = POSTGRES_PORT)
            conn.autocommit = True
            
            return conn
            
        except Exception as err:
            logging.error('Unable to connect to postgres db: %s', err)


def test_get_postgres_connection_unsuccess(monkeypatch):
        """ tests an exception is hit if the connection to postgres is unsuccesful"""
        
        # setup
        class mock_conn:
            def __init__(self, dbname, user, password, host, port):
                raise ConnectionError('This fake connection did not work')

            autocommit = False

        monkeypatch.setattr(psycopg2, 'connect', mock_conn)

I'm not able to successfully raise the exception in my mock function. Anyone know what I'm doing wrong here?

EDIT: Cleaned up code a bit

Answer 1

No need to create your own mock class - use unittest.mock.MagicMock instead.

You can use a MagicMock instance to mock basically anything, including third-party functions. If you add the side_effect=Exception() argument, an Exception is raised when the mock is called.

Python even allows you to do this in a context manager (with ... statement), so that the mocked function gets "un-mocked" once the context manager block ends.

Minimal example:

def some_external_lib():  # this is psycopg2.connect in your example
    pass

def my_func():  # this is get_postgres_connection in your example
    try:
        some_external_lib()
    except Exception as e:
        print(f"Error found: {e}")


import unittest
from unittest import mock

class TestMyFunc(unittest.TestCase):
    def test_my_func_external_lib_raises_exception(self):
        with mock.patch('__main__.some_external_lib', side_effect=Exception("ERROR")):
            my_func()


# Running example - prints the error message
t = TestMyFunc()
t.test_my_func_external_lib_raises_exception()

Note that the test, as it is written, doesn't actually test anything right now. Looking at the body of your get_postgres_connection function, you probably want to test whether it returns None, and whether something gets written to the log file, given that the exernal library raised an Exception.

Answer 2

Here is how I fixed the test function using mock, as answered above

    def test_get_postgres_connection_unsuccess():
        """ tests an exception is hit if the connection to postgres is unsuccesful"""

        # assert
        with mock.patch('psycopg2.connect', side_effect=Exception('ERROR')):
            self.pg_db.get_postgres_connection()