如何用Pytest测试Python百万收益计算程序的多返回值函数？

我编写了一款Python程序，用于根据用户输入的选项A/B及数字，计算在标普500（S&P500）投资中攒够100万美元所需的时间或每月需投入金额。现需使用Pytest为程序中的profile、option_check等函数编写测试用例，同时有疑问：如何测试返回多个值的函数？

程序代码 #

import math
import sys
# Global variables

s = 1000000 # final amount
r = float(11.88/100) # convert s&p500 annual return to effective monthly rate
i = float(r / 12) # effective rate

def main():
    x = input('Do you want to:\n (A): calculate how long it will take to save your first million investing a fixed amount in the S&P500 monthly?\n OR \n (B): Calculate how much you should save in the S&P500 per month to reach 1 000 000 within your time period?\n Type A or B: ')

    letter = profile(x)
    result = option_check(letter)
    print(result)
    sys.exit()

 # do users want how long it will take or how much to save within a period
def profile(x):
    while True:
        try:
            letter = x.upper()
            if letter == 'A' or letter == 'B':
                return letter
            else:
                main()
        except ValueError:
            print('INVALID !  Please enter: A or B')

def option_check(f):
    if f == 'A':
        pay_periods, amount = option_a()
        years, months = time(pay_periods)
        result = f"It will take {years} years and {months} months investing {amount} in the S&P500 to reach 1000 000"
        return result
    if f == 'B':
        amount, time_frame = option_b()
        result = f'You will reach 1 000 000 in {time_frame} years by investing {amount} in the S&P500 Monthly'
        return result

def option_a():
    # error check for ints
    while True:
        try:
            # monthly amount
            p = int(input('How much do you plan on saving in the S&P500 per month?\n'))

            # calculate years
            top = 1 - (s * (1 - (1 + i)) / p)

            periods = math.log(top) / math.log(1 + i)

            return periods, p

        except:
            print('Enter valid number\n')

def time(x):
    years = math.floor(x / 12)
    months = math.floor(12 * (x / 12 - math.floor(years)))
    return years, months

def option_b():
    while True:
        try:
            time_frame = int(input('What time frame do you have to reached 1 000 000 in years?\n'))

            # calculationg p: the monthly payment needed
            periods = time_frame * 12
            top = s - s * (1 + i)
            bottom = 1 - (1 + i ) ** periods
            amount = round(top / bottom , 2)
            return amount, time_frame
        except:
            print('Please Enter VALID NUMBER...\n')

if __name__ == '__main__':
    main()

现有测试代码 #

from forecast import option_check, time, profile

def main():
    test_time()
    test_option_check()

def test_time():
    assert time(12) == 1

def test_option_check():
    assert option_check('?')

def test_profile():
    assert profile('A')

if __name__ == '__main__':
    main()

一、如何测试返回多个值的函数 #

Python中函数返回多个值本质是返回元组，测试时直接断言函数返回的元组与预期元组完全相等即可。比如time函数返回(years, months)，测试时写assert time(14) == (1, 2)就可以验证结果是否正确。

二、修正并完善Pytest测试用例 #

以下是针对各函数的完整测试代码，解决了原测试中的错误（比如test_time原断言错误、option_check需要模拟输入等）：

import pytest
import math
from forecast import profile, time, option_check, option_a, option_b
from forecast import s, i  # 导入全局变量用于计算验证

def test_profile():
    # 测试合法输入
    assert profile('a') == 'A'
    assert profile('B') == 'B'
    # 注意：原profile函数中输入无效字符会递归调用main，这会导致测试陷入循环，建议后续重构该逻辑
    # 比如将递归改为抛出异常或返回提示，这里暂时只测试合法输入场景

def test_time():
    # 测试刚好整1年的情况
    assert time(12) == (1, 0)
    # 测试1年2个月的情况
    assert time(14) == (1, 2)
    # 测试2年1个月的情况
    assert time(25) == (2, 1)
    # 测试不足1年的情况
    assert time(5) == (0, 5)

def test_option_a(monkeypatch):
    # 模拟用户输入每月投入1000美元
    monkeypatch.setattr('builtins.input', lambda _: '1000')
    periods, amount = option_a()
    # 验证返回的金额是否正确
    assert amount == 1000
    # 验证计算的期数是否符合公式
    expected_top = 1 - (s * (1 - (1 + i)) / 1000)
    expected_periods = math.log(expected_top) / math.log(1 + i)
    assert abs(periods - expected_periods) < 1e-6  # 浮点误差容忍

def test_option_b(monkeypatch):
    # 模拟用户输入时间框架为20年
    monkeypatch.setattr('builtins.input', lambda _: '20')
    amount, time_frame = option_b()
    # 验证返回的时间框架是否正确
    assert time_frame == 20
    # 验证计算的每月投入是否符合公式
    periods = 20 * 12
    expected_top = s - s * (1 + i)
    expected_bottom = 1 - (1 + i) ** periods
    expected_amount = round(expected_top / expected_bottom, 2)
    assert amount == expected_amount

def test_option_check_a(monkeypatch):
    # 模拟选项A的输入流程：先选A，再输入每月1000
    monkeypatch.setattr('builtins.input', lambda _: '1000')
    result = option_check('A')
    # 计算预期结果中的年和月
    periods, _ = option_a()
    years, months = time(periods)
    expected_result = f"It will take {years} years and {months} months investing 1000 in the S&P500 to reach 1000 000"
    assert result == expected_result

def test_option_check_b(monkeypatch):
    # 模拟选项B的输入流程：选B，再输入20年
    monkeypatch.setattr('builtins.input', lambda _: '20')
    result = option_check('B')
    amount, _ = option_b()
    expected_result = f'You will reach 1 000 000 in 20 years by investing {amount} in the S&P500 Monthly'
    assert result == expected_result

三、关键说明 #

1. 模拟用户输入：使用pytest的monkeypatch工具替换builtins.input，避免测试时需要手动输入，实现自动化测试。
2. 浮点值断言：涉及计算的浮点结果，使用abs(实际值 - 预期值) < 1e-6来容忍计算中的微小误差。
3. 原代码的可测试性优化建议：profile函数中输入无效字符时递归调用main，会导致测试陷入无限循环，建议改为抛出自定义异常或返回错误提示字符串，方便测试。

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