在气流中模拟 Postgres 响应

Question

我试图模拟气流中的 postgres 响应，但我做不到。 尝试了在网上找到的所有解决方案，但没有任何帮助。 在气流手册上，它仅提供使用 docker image db 的解决方案。

这是代码：

from airflow.models import TaskInstance
from airflow.providers.postgres.hooks.postgres import PostgresHook

def process(pg_conn_id: str, ti: TaskInstance, **kwargs) -> str:

    xcom_in = ti.xcom_pull(task_ids="ids")

    pg_conn = PostgresHook(pg_conn_id).get_conn()
    cursor = pg_conn.cursor()
    cursor.execute("SELECT * FROM stub")
    result = cursor.fetchone()[0]
    print(result)

    return result

这是经过尝试的测试之一：

from unittest import mock

import pytest
from airflow.models import TaskInstance
from airflow.providers.postgres.hooks.postgres import PostgresHook

def test_process(mocker):
    input = "The input data"
    output="The value in DB"
    ti = TaskInstance
    ti.xcom_pull = MagicMock(return_value=input)


    mocker.patch.object(
        PostgresHook,
        "get_conn",
        return_value = MagicMock(name="pg")
    )

    assert output == process("mock_db", ti))

问题是结果，而不是我收到的“输入”值

<MagicMock name='pg.cursor().fetchone().__getitem__()' id='2259571508176'>

Answer 1

这就是模拟的工作方式——它们只是记录对它们的实际调用。 然后是调用的连接/链接，如果您在返回的模拟上调用一个方法，它将连接到原始模拟（下一段中更多内容）。

您成功地将pg_conn pg_conn = PostgresHook(pg_conn_id).get_conn()中的 pg_conn 制作为魔术模拟。 然后你有这一行cursor = pg_conn.cursor()所以cursor等于pg.cursor() 。 然后你得到等于cursor.fetchone()[0]的结果，我们将它连接到前一个 mock 以获得pg.cursor().fetchone().__getitem__() - 因为[0]在幕后调用__getitem__ 。

因此，在您的情况下，您可以使用以下代码：

pg_mock = MagicMock(name="pg")
pg_mock.cursor().fetchone().__getitem__.return_value = "input"

mocker.patch.object(
        PostgresHook,
        "get_conn",
        return_value = pg_mock
    )

关于模拟语法的最后一点说明：您通常可以将.return_value替换为() ，因此pg_mock.cursor().fetchone().__getitem__.return_value = "input"与pg_mock.cursor.return_value.fetchone.return_value.__getitem__.return_value = "input"相同pg_mock.cursor.return_value.fetchone.return_value.__getitem__.return_value = "input" 。