如何模拟使用上下文管理器的类中方法的返回值

Question

我有一个数据库模块

database.py

，其中有一个

Database

类，如下所示：

from someDatabase import theClient

class DatabaseClient:

    def __init__(self):
        self.connection = None

    def __enter__(self):
        self.connection = self.database_connection()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if self.connection:
            self.connection.close()

    def database_connection(self):
        client = theClient(<connection params>)
        return client

    def database_query(self, table: str, query: str):
        response = self.connection.search(
            body = query,
            table = table
        )
        return response

然后在我的烧瓶应用程序中我有一条路线：

from app.utils.Database import Database
from app.queries import thequery
from flask import Blueprint, jsonify

api = Blueprint('the_api', __name__)


@api.route("/api/this/route", methods=["GET"])
def get_some_stuff(**kwargs):
    try:

        # get input params from request

        with Database() as db:
            response = db.database_query("some_table", thequery)

        # transform the response object

        return jsonify(response), 200
    except Exception as e:
        raise e

我想弄清楚的是如何修补

DatabaseClient

，或者更具体地说，如何修补

database_query()

方法以返回我拥有的一些样本数据。

我想将方法返回值设置为某种值，然后在测试中调用路由，如下所示。我试图只使用 pytest，而不是与单元测试混合和匹配（如果可能的话）？

import pytest

from app import app

from tests.sample_data.api_responses.get_some_stuff import (
    get_some_stuff_api_response
)

from tests.sample_data.query_responses.get_some_stuff import (
    get_some_stuff_query_response
)


@pytest.fixture
def client():
    with app.test_client() as client:
        yield client


def test_get_some_stuff(client):


    # What is the correct patch here to set the return value of Database.database_query?
    # patch so it's equal to +get_some_stuff_query_response+

    # Call the route
    response = client.get(
        f"/api/this/route",
    )

    assert response.status_code == 200
    data = response.json
    assert data == get_some_stuff_api_response

Answer 1

您需要创建一个模拟上下文管理器来返回您想要的内容。

这是一个简单的示例，使用

monkeypatch

对文件打开上下文管理器执行相同的操作：

import builtins


def unit_under_test():
    with open("file.txt", "r") as f:
        return f.read()


class MockOpen:
    def __init__(self, *args, **kwargs):
        pass

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass

    def read(self):
        return "Hello, World!"


def test_unit_under_test(monkeypatch):
    monkeypatch.setattr(builtins, "open", MockOpen)

    assert unit_under_test() == "Hello, World!"

MockOpen

实现了上下文管理器所需的东西（

__enter__

和

__exit__

）。

所以，这翻译成你的例子会是这样的：

import app.utils.Database
...

class MockDatabase:
    def __init__(self, *args, **kwargs):
        pass

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass

    def database_query(self, table, query):
        return "Whatever the response needs to be"


def test_unit_under_test(monkeypatch):
    monkeypatch.setattr(app.utils.Database, "Database", MockDatabase)

    assert response.status_code == 200
    data = response.json
    assert data == get_some_stuff_api_response