CODEWITHSUNDEEP
pythonmockingpython-unittestcontextmanager
scrollout
scrollout

Reputation: 577

How do I unit test a custom context manager?

I have the following function that yields a pyodbc.connect() object to use with a with statement:

from contextlib import contextmanager
import pyodbc

@contextmanager
def get_db_connection(conn_string, **kwargs):
    try:
        conn = pyodbc.connect(conn_string, **kwargs)
        yield conn
    except Exception as connection_error:
        raise ValueError('Could not connect to db.', connection_error) from None
    finally:
        conn.close()

I'd like to write unit tests for this function to test that a connection can be opened, closed, and when an error occurs, a ValueError is raised. I have:

from unittest.mock import Mock, patch
from db.query import get_db_connection

@patch('db.query.pyodbc.connect')
def test_get_db_connection(self, mock_connect):
    conn_string = Mock()
    with get_db_connection(conn_string) as conn:
        pass
    # print(conn) is acceptable here and prints:
    # <MagicMock name='connect()' id='2595126451264'>
    mock_connect.assert_called_once()

    mock_connect.side_effect = Exception()
    def _t():
        with get_db_connection(conn_string) as conn:
            pass

    self.assertRaises(UnboundLocalError, _t)

Everything here runs and passes, but I've noticed the conn is available after the with statement scope should be closed.

How do I write a unit test for get_db_connection() using mock so that the conn variables behave as I'd expect them? Or is there another way to sufficiently unit test it?

Upvotes: 3

Views: 3149

Answers (1)

scrollout
scrollout

Reputation: 577

Use

  • conn.close.assert_not_called() to check that the connection is opened in the first with block.
  • conn.close.assert_called_once() to check that the connection is closed after the first with block.
  • There is already a check for no connection being established when an error occurs in the try block - UnboundLocalError: local variable 'conn' referenced before assignment

Upvotes: 2

Related Questions