如何对自定义上下文管理器进行单元测试?

How do I unit test a custom context manager?

我有以下函数生成一个 pyodbc.connect() 对象与 with 语句一起使用:

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()

我想为此函数编写单元测试,以测试连接是否可以打开、关闭,以及在发生错误时引发 ValueError。我有:

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)

这里的一切都运行并通过,但我注意到 connwith 语句范围应该关闭后可用。

如何使用 mockget_db_connection() 编写单元测试,以便 conn 变量的行为符合我的预期?还是有另一种方法可以对其进行充分的单元测试?

使用

  • conn.close.assert_not_called() 检查连接是否在第一个 with 块中打开。
  • conn.close.assert_called_once() 检查第一个 with 块后连接是否关闭。
  • try 块中发生错误时,已经检查是否未建立连接 - UnboundLocalError:局部变量 'conn' 在赋值前被引用