时间戳

首页 > 解决方案 > 如何对自定义上下文管理器进行单元测试？

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

问题描述

我有以下函数，它产生一个与语句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引发 a。我有：

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)

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

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

标签： pythonmockingpython-unittestcontextmanager

解决方案

利用

  • conn.close.assert_not_called()检查连接是否在第一个with块中打开。
  • conn.close.assert_called_once()检查连接是否在第一个with块之后关闭。
  • 当块中发生错误时，已经检查了是否没有建立连接try- UnboundLocalError: local variable 'conn' referenced before assignment

推荐阅读