17370845950

本文介绍在不修改函数返回值、不依赖文件系统的情况下，通过 `unittest.mock` 捕获被测函数内部创建的 pandas dataframe，核心是利用 `side_effect` 或 `wraps` 机制访问原始实例。

在单元测试中验证 DataFrame 的最终状态（如列名、形状、数值）是一项常见需求，但当目标函数设计为“仅执行副作用”（例如调用 df.to_csv() 后无返回值），直接获取其内部 DataFrame 实例就成为挑战。由于 Airflow 等环境限制不能返回 DataFrame，且避免读写真实文件，我们需要一种零侵入、纯内存、可断言的捕获方式。

关键思路是：不只 mock to_csv 的行为，更要让它“泄露”调用它的 DataFrame 实例本身。unittest.mock 提供了两种可靠方案：

✅ 方案一：使用 wraps 保留原始方法，并在调用时捕获 self

# test.py
import unittest
from unittest.mock import patch
from foo import bar
import pandas as pd

class TestBar(unittest.TestCase):
    def test_bar_captures_dataframe(self):
        captured_df = None

        def capture_df_instance(self, *args, **kwargs):
            nonlocal captured_df
            captured_df = self  # ← self 就是调用 to_csv 的 DataFrame 实例
            return None  # 保持原 to_csv 行为（不真正写文件）

        with patch("pandas.DataFrame.to_csv", wraps=capture_df_instance):
            bar()

        # 断言捕获到的 DataFrame
        self.assertIsInstance(captured_df, pd.DataFrame)
        self.assertEqual(list(captured_df.columns), ['a', 'b'])
        self.assertEqual(captured_df.shape, (3, 2))
        self.assertEqual(captured_df['a'].tolist(), [1, 2, 3])

✅ 优势：语义清晰，self 即 DataFrame；无需修改被测代码；完全绕过文件 I/O。

✅ 方案二：使用 side_effect + mock_calls 反向提取（适用于更复杂场景）

@patch("pandas.DataFrame.to_csv")
def test_bar_with_side_effect(self, mock_to_csv):
    # 让 mock 记录每次调用的 self（即 DataFrame）
    mock_to_csv.side_effect = lambda self, *args, **kwargs: setattr(self, '_test_captured', True)

    bar()

    # 从 mock 调用中提取第一个调用的 self（需确保只调用一次）
    if mock_to_csv.called:
        df = mock_to_csv.call_args[0][0]  # ← call_args[0][0] 是第一个位置参数，即 self
        self.assertIsInstance(df, pd.DataFrame)
        # 后续断言同上...

⚠️ 注意事项：

• wraps 方案更推荐：它不破坏原始方法签名，且 self 引用稳定可靠；
• 避免过度依赖 mock_calls 解析——call_args 的索引易受参数变化影响；
• 若函数中 to_csv 被多次调用，需结合 call_count 或 assert_has_calls 精确匹配；
• 确保 pandas 在测试前已导入（import pandas as pd），否则 wraps 可能因未解析类型而失败。

总结：测试驱动开发（TDD）并不强制重构业务逻辑，但要求可观测性。通过 wraps 或 side_effect “钩住”方法调用链中的实例，我们能在零副作用前提下完成深度断言——这既是技术技巧，也是对“可测试性即设计质量”的践行。