Python издевается над psycopg2 соединением и курсором

Question

Мне не удалось смоделировать соединение с базой данных 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 показывают, что базовый объект не является волшебным макетомКак можно издеваться над глобальным соединением БД и курсором БД, используя менеджер контекста? Любая помощь приветствуется.

Answer 1

Я нашел решение! У меня было чувство, что мои проблемы связаны с глобальной переменной 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

Ничто не сравнится с ответом на ваш собственный вопрос!