使用模拟获取应用的函数输入数据框

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. 最后一个函数不返回数据帧（它只是迭代所有行，类似于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.2Pandas 版本，升级到 Pandas1.0.5不会改变问题。

我尝试在熊猫问题中进行搜索，但没有找到有关此主题的任何信息。

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()