Python unittest如何访问被测函数内部变量编写相等断言

函数内部的局部变量默认会在函数执行结束后随栈帧销毁，直接通过导入调用的方式拿不到是正常的，针对你的需求有三种可落地的实现方案，按推荐优先级排序：

方案1：优先测业务行为，完全不需要碰内部变量（最推荐） #

你要校验的核心业务规则是「输入数据列数和处理后输出列数相等」，根本不需要掏函数内部的colinput、coloutput变量——这俩值本质就是输入、输出DataFrame的列数，你自己构造可控的测试输入，直接统计输入输出的列数做断言就行，完全不侵入原业务代码。
示例代码：

import unittest
import pandas as pd
from mainscript import getweight

class TestGetweight(unittest.TestCase):
    def test_col_count_match(self):
        # 构造列数明确的测试输入
        test_df = pd.DataFrame({
            "feat1": [1,2,3],
            "feat2": [4,5,6],
            "feat3": [7,8,9]
        })
        input_col_num = len(test_df.columns)
        # 执行函数拿到输出
        output_df = getweight(test_df)
        output_col_num = len(output_df.columns)
        # 断言符合业务要求
        self.assertEqual(input_col_num, output_col_num)

这种写法属于标准的黑盒测试，和函数内部实现完全解耦：后续哪怕你重构getweight、把内部两个临时变量改名甚至删掉，只要「输入输出列数相等」的业务规则没变，测试用例就不会失效，维护成本极低。

方案2：必须校验内部变量值时，用栈追踪捕获局部变量 #

如果你的需求是必须确认函数内部统计的colinput、coloutput本身计算逻辑没写错（比如怕内部统计列数的时候把索引列算进去之类的场景），可以用Python的系统追踪钩子，在函数返回前抓出栈帧里的局部变量，不需要修改原函数代码。
示例代码：

import unittest
import sys
import pandas as pd
from mainscript import getweight

class TestGetweight(unittest.TestCase):
    def test_internal_col_var_equal(self):
        captured = {}
        # 定义追踪函数，捕获getweight返回时的局部变量
        def tracer(frame, event, arg):
            if event == "return" and frame.f_code.co_name == "getweight":
                captured.update(frame.f_locals)
            return tracer
        
        test_df = pd.DataFrame({"a": [1], "b": [2]})
        # 开启追踪执行函数
        sys.settrace(tracer)
        getweight(test_df)
        sys.settrace(None)

        # 直接断言捕获到的内部变量
        self.assertEqual(captured["colinput"], captured["coloutput"])

注意：这个方案属于强耦合的白盒测试，只要你后续修改函数内部的变量名、调整变量作用域，测试就会直接报错，非必要不要用。

方案3：给函数加调试参数（长期维护最方便） #

如果这两个统计值后续会被频繁校验，可以给getweight加一个可选的调试开关，不影响原有业务调用的前提下，需要的时候把内部统计值返回：
原函数修改示例：

# mainscript 中的getweight调整
def getweight(input_df, debug=False):
    colinput = len(input_df.columns)
    # 原有数据处理逻辑
    # ... 处理完成得到 output_df
    coloutput = len(output_df.columns)
    
    if debug:
        return output_df, {"colinput": colinput, "coloutput": coloutput}
    return output_df

测试时传入debug=True就能直接拿到两个内部值，业务侧正常调用不需要传参，完全不受影响。

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