通过模拟获取应用程序的函数输入数据框

Question

我有以下功能

def main():
    (
        pd.DataFrame({'a': [1, 2, float('NaN')], 'b': [1.0, 2, 3]})
        .dropna(subset=['a'])
        .assign(
            b=lambda x: x['b'] * 2
        )
        .apply(do_something_with_each_row, axis='columns')
    )

def do_something_with_each_row(one_row):
    # do_something_with_row
    print(one_row)

在我的测试中，我想查看在所有链式操作之后构建的数据帧，并在调用 do_something_with_each_row 之前检查它是否一切正常。最后一个函数不是 return 数据框（它只是迭代所有行，类似于 iterrow）。

我试过像这样模拟 apply 函数：

# need pytest-mock and pytest
import pandas as pd


def test_not_working(mocker):
    mocked_apply = mocker.patch.object(pd.Dataframe, 'apply')
    main()

但在这种情况下，我无法访问输入到 apply 的数据框以测试其内容。

我也试过模拟 do_something_with_each_row:

# need pytest-mock and pytest
import pandas as pd


def test_not_working_again(mocker):
    mocked_to_something = mocker.patch('path.to.file.do_something_with_each_row')
    main()

但这次我有所有带有行参数的调用，但它们都有 None 个值。

如何获取调用 apply 函数的数据帧并检查它是否确实与以下内容相同：

pd.Dataframe({'a': [1, 2], 'b': [2.0, 4]})

我正在使用 0.24.2 pandas 版本，升级到 pandas 1.0.5 不会改变问题。

我尝试搜索 pandas 个问题，但没有找到任何关于这个主题的信息。

Answer 1

如果我正确理解了你的问题，这是获得你想要的行为的方法之一：

def test_i_think_this_is_what_you_asked(mocker):
    original_apply = pd.DataFrame.apply
    def mocked_apply(self, *args, **kw):
        assert len(self) == 2 # self is the pd.DataFrame at the time apply is called
        assert self.a[0] == 1
        assert self.a[1] == 3 # this will fail cause the value is 2
        assert self.b[0] == 2.0
        assert self.b[1] == 4.0
        return original_apply(self, *args, **kw)
    mocker.patch.object(pd.DataFrame, 'apply', side_effect=mocked_apply, autospec=True)
    main()

通过模拟获取应用程序的函数输入数据框

Get apply's function input dataframe with mocking

pytest

python-3.x

pandas

pytest-mock