如何在pytest的VRP测试预期结果中使用逻辑运算符？

问题背景 #

用pytest测试路由优化（VRP）代码时，自定义了VRPResults数据类存储结果，通过__eq__方法做相等校验，但仅支持字段完全匹配，无法满足诸如「总旅行时间小于20分钟」这类非相等的校验场景。

解决方案 #

方案1：扩展VRPResults类支持自定义校验逻辑 #

给VRPResults新增校验规则字段，允许为每个字段指定自定义校验逻辑，替代原有的__eq__做更灵活的校验：

from dataclasses import dataclass, fields
from typing import Optional, List, Dict, Callable

@dataclass(slots=True)
class VRPResults:
    solver_time: Optional[float] = None
    total_travel_time: Optional[int] = None
    route_lengths: Optional[List[int]] = None
    route_indices: Optional[List[List[int]]] = None
    # 存储字段校验规则：键为字段名，值为接收实际值并返回bool的校验函数
    validation_rules: Optional[Dict[str, Callable[[any], bool]]] = None

    def __post_init__(self):
        if self.route_lengths is not None and self.total_travel_time is None:
            self.total_travel_time = sum(self.route_lengths)
        
        # 初始化默认规则：未指定规则时，对非None的预期字段做相等校验
        if self.validation_rules is None:
            self.validation_rules = {}
            for field in fields(VRPResults):
                if field.name == "validation_rules":
                    continue
                expected_val = getattr(self, field.name)
                if expected_val is not None:
                    self.validation_rules[field.name] = lambda x, exp=expected_val: x == exp

    def validate(self, other: "VRPResults") -> bool:
        if not isinstance(other, VRPResults):
            raise TypeError(f"无法校验 {type(self)} 和 {type(other)} 类型的对象")
        
        for field_name, rule in self.validation_rules.items():
            actual_val = getattr(other, field_name)
            if actual_val is None or not rule(actual_val):
                return False
        return True

使用示例

• 场景1：总旅行时间等于10

  vrp_results_expected = VRPResults(total_travel_time=10)
  assert vrp_results_expected.validate(vrp_results_predicted), "结果不符合预期"
  

• 场景3：总旅行时间小于20

  vrp_results_expected = VRPResults(
      validation_rules={
          "total_travel_time": lambda x: x < 20
      }
  )
  assert vrp_results_expected.validate(vrp_results_predicted), "总旅行时间超出阈值"
  

• 场景2：同时校验路径索引和长度相等

  vrp_results_expected = VRPResults(
      route_indices=[[0,1,2]],
      route_lengths=[5,2,4]
  )
  assert vrp_results_expected.validate(vrp_results_predicted), "路径信息不符合预期"
  

方案2：使用pytest现成工具或自定义断言函数 #

用pytest-check插件实现多类型断言

pytest-check支持多种断言方式，无需修改数据类，直接在测试用例中编写灵活校验：

1. 安装插件：

   pip install pytest-check
   

2. 测试用例示例：

   import pytest_check as check
   
   def test_vrp_scenario3():
       vrp_data, _ = create_data_model()
       vrp_results_predicted = solve_vrp(vrp_data)
       # 校验总旅行时间小于20
       check.less(vrp_results_predicted.total_travel_time, 20)
       # 额外校验求解时间小于10秒
       check.less(vrp_results_predicted.solver_time, 10.0)
   

自定义断言函数

如果不想引入第三方插件，可以自己封装断言函数，支持运算符+阈值的组合：

def assert_vrp_results(actual: VRPResults, **kwargs):
    """自定义VRP结果断言，支持指定字段的校验规则"""
    for field_name, expected in kwargs.items():
        actual_val = getattr(actual, field_name)
        if isinstance(expected, tuple) and len(expected) == 2:
            op, threshold = expected
            if op == "lt":
                assert actual_val < threshold, f"{field_name} 应小于 {threshold}，实际为 {actual_val}"
            elif op == "le":
                assert actual_val <= threshold, f"{field_name} 应小于等于 {threshold}，实际为 {actual_val}"
            elif op == "gt":
                assert actual_val > threshold, f"{field_name} 应大于 {threshold}，实际为 {actual_val}"
            elif op == "ge":
                assert actual_val >= threshold, f"{field_name} 应大于等于 {threshold}，实际为 {actual_val}"
        else:
            # 默认做相等校验
            assert actual_val == expected, f"{field_name} 应为 {expected}，实际为 {actual_val}"

使用示例

def test_vrp_scenario3():
    vrp_data, _ = create_data_model()
    vrp_results_predicted = solve_vrp(vrp_data)
    assert_vrp_results(vrp_results_predicted, total_travel_time=("lt", 20))

def test_vrp_scenario1():
    vrp_data, _ = create_data_model()
    vrp_results_predicted = solve_vrp(vrp_data)
    assert_vrp_results(vrp_results_predicted, total_travel_time=10)

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