Python含用户输入、循环、函数依赖的ATM程序单元测试方案

问题背景 #

我写了一个实现基础功能的ATM Python程序：

• 校验用户输入的PIN码
• PIN匹配后允许用户输入取款金额
• 从账户余额中扣除对应取款金额
  目前代码耦合度很高：函数之间直接链式调用，前一个函数的返回值直接作为后一个函数的入参，根本没法单独给单个函数写单元测试——只要尝试单独运行某个函数，就会触发整个程序跑全流程，依次要求输入各个环节的内容。再加上get_pin和get_amount里直接写死了处理用户输入的while循环，进一步加大了写单测的难度。

待确认问题 #

这种场景下有没有不重构就能写单元测试的可行方法？还是说最优解就是重构代码，把用户输入、while循环逻辑拆出去，让每个独立函数只靠传入参数处理核心逻辑？

需要覆盖的测试点 #

• withdraw_cash函数：验证取款金额扣减逻辑是否正确
• get_pin函数：验证PIN码匹配判断逻辑是否正常
• get_amount函数：验证输入非法金额时是否正常触发错误提示

原始代码 #

user = {
    'pin': 1234
}

def withdraw_cash(amount):
    balance_account = 100
    if amount > balance_account:
        raise ValueError("You don't have sufficient balance to make this withdrawal")
    else:
        new_balance = balance_account - amount
        return new_balance
    
    
def get_pin():
    count = 0
    to_exit = False
    while (count < 3) and (not to_exit):
        try:
            pin = int(input('Please enter your four digit pin: '))
        except ValueError:
            print("Please enter correct pin")
            count += 1
        if pin != user['pin']:
            print("Pin does not match.. Try Again")
            count += 1
        else:
            return get_amount(pin)        
    if count == 3:
        a = '3 UNSUCCESFUL PIN ATTEMPTS, EXITING \n !!!!!YOUR CARD HAS BEEN LOCKED!!!!!'
        return a 
            


def get_amount(pin):
    while True:
        try:
            amount = int(input("Enter the amount of money you want to withdraw: "))
        except ValueError as v:
            print(f"Enter correct amount: ")
        else:
            return withdraw_cash(amount)

try:
    get_pin()
except ValueError as v:
    print(f"ERROR: {v}")

回答 #

两种路径都能实现测试覆盖，但长期来看优先选重构方案，mock打桩的方式只适合临时用。

首先withdraw_cash本身现在就可以直接测，它没有依赖任何输入输出逻辑，只是硬编码了初始余额100，测试的时候直接传参即可：传50应该返回50，传超过100的数值应该抛出余额不足的错误，这部分没有任何测试障碍。

真正难测的是get_pin和get_amount，两个函数同时绑定了输入采集、循环控制、下游业务调用，才会导致一运行就触发全流程。

不重构的临时测试方案 #

如果暂时不想改代码，可以用Python标准库unittest.mock做打桩测试：

• 测试get_pin时，把builtins.input替换成预设返回序列的mock对象，按测试用例需求依次返回非数字值、错误PIN、正确PIN，同时把下游的get_amount也mock掉阻止实际执行，就可以分别验证：非数字输入是否正常计数、错误PIN是否弹出提示、连续3次错误是否返回锁卡提示、正确PIN是否触发后续流程。
• 测试get_amount时同理，mockinput依次返回非数字值、超额取款值、合法取款值，同时把下游的withdraw_cashmock掉，就能验证非法输入是否弹出对应错误提示、合法输入是否调用扣款逻辑。

这种方案的问题非常明显：测试代码和实现细节强绑定，只要你调整了input调用顺序、循环内的判断逻辑，测试直接就会失败，本质是测试在“测代码怎么写”而不是“测逻辑对不对”，长期维护成本很高。

推荐方案：轻量重构拆分逻辑 #

只需要花十几分钟调整代码结构，就能把测试成本降到最低，后续改代码也更方便：

• 把所有核心判断逻辑抽成无IO依赖的纯函数：比如单独写PIN校验函数，传入用户输入的PIN和正确PIN，直接返回校验结果；单独写金额校验函数，传入用户输入的内容和账户余额，返回是否合法、错误提示、转换后的数值；把withdraw_cash改成接收余额、取款金额两个参数，去掉硬编码的余额值。这些纯函数只靠入参计算返回值，写单测的时候连mock都不需要，直接传参断言结果即可，稳定又好写。
• 把while循环、input调用、函数串联的逻辑全部放到最外层的主流程函数里，这部分本身的职责就是做交互控制、流程跳转，不需要做细粒度的单元测试，端到端测试的时候覆盖即可。

拆分完之后你需要的三个测试点全部落在纯逻辑函数上，单测写起来非常简单，后续改业务逻辑也不会因为调整交互流程导致测试大面积失效。

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