如何从Mock的to_csv方法中获取DataFrame实例用于测试？

你可以通过Mock的side_effect功能捕获调用to_csv的DataFrame实例，具体实现如下：

1. 在测试函数中定义一个捕获逻辑，接收self（即调用to_csv的DataFrame实例）和其他参数，将self保存到外部变量中。
2. 将该捕获逻辑赋值给Mock对象的side_effect属性，当to_csv被调用时会自动执行捕获逻辑，从而拿到目标DataFrame。

修改后的测试代码如下：

import unittest
from unittest.mock import patch
import pandas as pd
from foo import bar

class TestBar(unittest.TestCase):
    @patch("pandas.DataFrame.to_csv")
    def test_bar(self, mock_to_csv):        
        captured_df = None
        
        def capture_to_csv(self, *args, **kwargs):
            nonlocal captured_df
            captured_df = self
        
        mock_to_csv.side_effect = capture_to_csv
        
        bar()
        
        # 验证捕获到的DataFrame是否符合预期
        expected_df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
        pd.testing.assert_frame_equal(captured_df, expected_df)

这个方案无需修改原函数（满足Airflow的限制），也不需要读写文件（避开测试环境文件系统不可预测的问题），直接就能获取函数内部处理后的最终DataFrame实例。

内容的提问来源于stack exchange，提问作者Dracula