Python模拟Psycopg2连接和游标

时间:2019-10-28 21:47:23

标签: python unit-testing mocking pytest contextmanager

由于无法重构psycopg2数据库连接和游标,因此无法对其进行模拟,因为我将其重构为对游标使用上下文管理器。我知道,使用上下文管理器,还有一些其他的魔术方法可用于资源设置和清理(__enter____exit__),但是即使将其添加到组合中也无法缩小问题的范围。 。

代码如下:

import os
import psycopg2
import psycopg2.extras

DB_HOST = os.getenv('DB_HOST')     
DB_PORT = os.getenv('DB_PORT')
DB_NAME = os.getenv('DB_NAME')
DB_USER = os.getenv('DB_USER')
DB_PASSWORD = os.getenv('DB_PASSWORD')            
CONN = psycopg2.connect(f'dbname={DB_NAME} user={DB_USER} host={DB_HOST} port={DB_PORT} password={DB_PASSWORD}')

def my_func():
    message = None
    print(CONN)  # Added to debug - this never prints out a magic mock reference

    select_sql = 'SELECT * FROM app.users WHERE name = %s LIMIT 1;'
    with CONN.cursor(cursor_factory = psycopg2.extras.DictCursor) as cursor:
        print(cursor) # Debug - no magic mock reference
        cursor.execute(select_sql, ("Bob"))
        row = cursor.fetchone()

    if row is not None: 
        message = "name found"
    else:
        message = "name not found"

    return message

这是测试代码,也是我尝试模拟连接和游标的尝试

import pytest

from src import my_class
from unittest import mock

class TestFunction:
    @mock.patch('psycopg2.connect')
    def test_my_func(self, mock_connect):

        mock_cursor = mock.MagicMock()
        mock_cursor.__enter__.return_value.fetchone.return_value = {
            "id": 1,
            "name": "Bob",
            "age": 25
        }
        mock_connect.return_value.cursor.return_value = mock_cursor

        result = my_class.my_func()
        assert result == "found"

如果我在全局连接cursor上调用CONN,则不确定如何模拟连接或游标。当前正在运行的pytest显示测试失败,并且print语句显示基础对象不是魔术模型。如何使用上下文管理器模拟全局数据库连接和数据库游标?任何帮助表示赞赏。

1 个答案:

答案 0 :(得分:0)

我找到了解决方案!我感到我的问题出在全局变量CONN周围,并研究了如何模拟python全局变量并进行了尝试。我的印刷声明最终表明,这些实例确实是魔术嘲讽。

有意义的是,在我的测试中,我仅调用被测方法,而不用全局变量触发代码-因此只需将其设置为所需的值即可。找到答案后听起来很直接,对吧?

mock_cursor = mock.MagicMock()
mock_cursor.fetchone.return_value = {
      "id": 1,
      "name": "Bob",
      "age": 25
}
mock_connect.cursor.return_value.__enter__.return_value = mock_cursor
my_class.CONN = mock_connect

没什么比回答自己的问题了!