仅模拟类的一个函数

Question

我想编写一个单元测试，可以确保函数调用中的 SQL 语句在原理上是正确的。它应该测试此调用的执行。然后我想模拟对提交的调用，以便不会发生对数据库的插入。我正在使用 psycopg2 进行测试。

我有一个类似的功能：

def test_insert(a, b, c):
    con = psycopg2.connect(os.environ['PGDB'])
    cur = con.cursor()
    cur.execute('insert into test_table values ({a}, {b}, {c})'.format(a=a, b=b, c=c))
    con.commit()
    con.close()

调用

test_insert(1,2,3)

时，我看到该行插入到表中。现在我尝试模拟这个电话。到目前为止我已经采取了一些方法：

@mock.patch('psycopg2.connect')
def test(mock_connect, a, b, c):
    mock_con = mock_connect.return_value
    mock_con.commit.return_value = None
    insert_row(a, b, c)

这看似可行，但实际上并没有调用执行语句。例如

test_insert(1,4,'xyz')

会失败，而

test(1,4,'xyz')

则不会。接下来我尝试仅模拟 psycopg2 中连接类的提交方法：

@mock.patch('psycopg2.extensions.connection.commit')
def test_insert(mock_commit, a, b, c):
    mock_commit.return_value = None
    insert_row(a,b,c)

但这给了我一个语法错误

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/a/.virtualenv/test/lib/python2.7/site-packages/mock/mock.py", line 1318, in patched
    patching.__exit__(*exc_info)
  File "/home/a/.virtualenv/test/lib/python2.7/site-packages/mock/mock.py", line 1480, in __exit__
    setattr(self.target, self.attribute, self.temp_original)
TypeError: can't set attributes of built-in/extension type 'psycopg2.extensions.connection'

有什么好方法来做我想做的事吗？

Answer 1

我假设您正在使用 pytest 并且如果它们不是实际测试，那么以

test_

开头命名您的函数并不是一个好习惯，因为这可能会引发测试框架的问题。因此，我稍微修改了您的初始代码片段，如下所示，并将该模块命名为

psyco.py

import psycopg2
import os

def insert(a, b, c):
    con = psycopg2.connect(os.environ['PGDB'])
    import ipdb; ipdb.set_trace()
    cur = con.cursor()
    cur.execute('insert into test_table values ({a}, {b}, {c})'.format(a=a, b=b, c=c))
    con.commit()
    con.close()

接下来，我通过考虑补丁如何工作和在哪里打补丁为您的方法创建了测试。当您处理操作系统环境变量时，这个问题可以帮助您理解为什么我这样嘲笑它。

测试的示例实现如下：

from psyco import insert
from unittest.mock import patch, Mock, MagicMock
import os

@patch.dict(os.environ,{'PGDB':'db_url'})
@patch('psycopg2.connect')
def test_insert_function(psycopg2_mock):

    x = 1
    y = 4
    z = 'xyz'
    sql_query = 'insert into test_table values ({0}, {1}, {2})'.format(x,y,z)
    insert(x,y,z)
    assert psycopg2_mock.return_value.cursor.call_count == 1
    psycopg2_mock.return_value.cursor.return_value.execute.assert_called_with(sql_query)
    assert psycopg2_mock.return_value.commit.call_count == 1
    assert psycopg2_mock.return_value.close.call_count == 1

Answer 2

这篇文章可能会有所帮助，在 Python 中使模拟函数强类型化