如何为使用random.choice的Python老虎机函数编写单元测试

因为函数依赖random.choice的随机性，直接测试会每次结果不同，核心解决思路是Mock随机模块——固定random.choice的返回值，从而验证函数的逻辑正确性（比如行列数是否符合要求、符号来源是否正确、列内是否无重复选取等）。

用unittest框架的测试示例 #

假设函数放在slot_machine.py文件中，测试代码如下：

1. 导入依赖 #

import unittest
from unittest.mock import patch
import random
from slot_machine import get_slot_machine_spin

2. 编写测试用例 #

class TestSlotMachineSpin(unittest.TestCase):
    def test_spin_dimensions(self):
        """验证返回的行列数与输入参数一致"""
        rows = 3
        cols = 5
        symbols = {"A": 2, "B": 3, "C": 4}
        
        # 固定random.choice的返回序列，覆盖所有行列的选取需求
        with patch('random.choice') as mock_choice:
            mock_choice.side_effect = ["A", "B", "C"] * cols
            
            result = get_slot_machine_spin(rows, cols, symbols)
            
            # 检查列数是否等于cols
            self.assertEqual(len(result), cols)
            # 检查每一列的行数是否等于rows
            for column in result:
                self.assertEqual(len(column), rows)
    
    def test_spin_symbol_rules(self):
        """验证返回符号来自输入集合，且单列内无重复"""
        rows = 2
        cols = 2
        symbols = {"X": 2, "Y": 2}
        
        with patch('random.choice') as mock_choice:
            # 构造符合"单列无重复"的返回序列
            mock_choice.side_effect = ["X", "Y", "Y", "X"]
            
            result = get_slot_machine_spin(rows, cols, symbols)
            
            # 检查所有符号都来自输入的symbols
            for column in result:
                for symbol in column:
                    self.assertIn(symbol, symbols.keys())
            
            # 检查单列内没有重复符号（因为每次选取后会从current_symbols中移除）
            for column in result:
                self.assertEqual(len(column), len(set(column)))

if __name__ == '__main__':
    unittest.main()

用pytest的简化版本 #

如果用pytest框架搭配pytest-mock插件，代码会更简洁：

import pytest
import random
from slot_machine import get_slot_machine_spin

def test_spin_dimensions(mocker):
    rows = 3
    cols = 5
    symbols = {"A":2, "B":3, "C":4}
    
    mocker.patch('random.choice', side_effect=["A", "B", "C"]*cols)
    
    result = get_slot_machine_spin(rows, cols, symbols)
    
    assert len(result) == cols
    assert all(len(col) == rows for col in result)

def test_spin_no_duplicates_in_column(mocker):
    rows = 2
    cols = 2
    symbols = {"X":2, "Y":2}
    
    mocker.patch('random.choice', side_effect=["X", "Y", "Y", "X"])
    
    result = get_slot_machine_spin(rows, cols, symbols)
    
    for col in result:
        assert len(set(col)) == len(col)  # 单列无重复
        assert all(s in symbols for s in col)  # 符号均来自输入

额外优化：拆分逻辑提升可测试性 #

原函数中构造all_symbols的逻辑可以拆成单独的辅助函数，方便单独验证符号列表的正确性：

def build_all_symbols(symbols):
    all_symbols = []
    for symbol, count in symbols.items():
        all_symbols.extend([symbol] * count)
    return all_symbols

def get_slot_machine_spin(rows, cols, symbols):
    all_symbols = build_all_symbols(symbols)
    columns = []
    for _ in range(cols):
        column = []
        current_symbols = all_symbols[:] 
        for _ in range(rows):
            value = random.choice(current_symbols)
            current_symbols.remove(value)
            column.append(value)
        columns.append(column)
    return columns

这样可以单独测试build_all_symbols函数，确保符号列表构造符合预期，再聚焦测试随机选取的逻辑。

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