Locust子类继承基类传参触发KeyError问题求助

问题场景 #

想要通过继承BaseUser基类实现多子类复用核心逻辑，仅让子类差异体现在dto_class和api_suffix两个参数上，但运行脚本时触发KeyError: 'dto_class'。

基类代码（locust_base.py） #

### locust_base.py ###
from locust import HttpUser, between, task
## 其他导入 ##

class BaseUser(HttpUser):
    """
    通用逻辑实现
    """
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dto_class = kwargs.pop('dto_class')
        self.api_suffix: str = kwargs.pop('api_suffix')
        """
        其他初始化操作
        """

子类代码（locust_specific_dto.py） #

"""
locust_specific_dto.py
"""
from locust_base import BaseUser

class SpecificDtoUser(BaseUser):
    def __init__(self, *args, **kwargs):
        super(SpecificDtoUser, self).__init__(dto_class=SpecificDto, api_suffix="specificEndpoint", *args, **kwargs)

错误信息 #

KeyError: 'dto_class'
2024-05-09T18:09:14Z <Greenlet at 0x21e4d99d580: <lambda>> failed with KeyError

[2024-05-09 14:09:14,811] WSAMZN-SPLR8QE1/CRITICAL/locust.runners: Unhandled exception in greenlet: <Greenlet at 0x21e4d99d580: <lambda>>
Traceback (most recent call last):
  File "src\gevent\greenlet.py", line 908, in gevent._gevent_cgreenlet.Greenlet.run
  File "{REDACTED}\AppData\Roaming\Python\Python311\site-packages\locust\runners.py", line 553, in <lambda>
    lambda: self._start(user_count, spawn_rate, wait=wait, user_classes=user_classes)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "{REDACTED}\AppData\Roaming\Python\Python311\site-packages\locust\runners.py", line 524, in _start
    self.spawn_users(user_classes_spawn_count, wait)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "{REDACTED}\AppData\Roaming\Python\Python311\site-packages\locust\runners.py", line 238, in spawn_users
    new_users += spawn(user_class, spawn_count)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "{REDACTED}\AppData\Roaming\Python\Python311\site-packages\locust\runners.py", line 227, in spawn
    new_user = self.user_classes_by_name[user_class](self.environment)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "{REDACTED}\Documents\dev\{REDACTED}\locust_base.py", line 23, in __init__
    self.dto_class = kwargs.pop('dto_class')
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyError: 'dto_class'

问题原因 #

1. Locust实例化User类时，只会传入environment这一个参数，不会传递自定义关键字参数，导致父类__init__中kwargs.pop('dto_class')找不到对应键值。
2. 子类调用父类构造方法时参数顺序错误：Python语法要求*args必须放在关键字参数之前，原代码将关键字参数放在*args前面，导致参数无法被正确传入kwargs。

解决方案 #

放弃通过构造方法传递参数，将dto_class和api_suffix定义为类属性，子类直接覆写这些属性即可，无需重写__init__方法，既符合Locust设计逻辑，也能实现逻辑复用。

修改后的基类代码 #

### locust_base.py ###
from locust import HttpUser, between, task
## 其他导入 ##

class BaseUser(HttpUser):
    """
    通用逻辑实现
    """
    # 定义默认类属性，子类可覆写
    dto_class = None
    api_suffix = ""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 可选：添加校验确保子类必须覆写属性
        if self.dto_class is None:
            raise ValueError("子类必须设置dto_class属性")
        if not self.api_suffix:
            raise ValueError("子类必须设置api_suffix属性")
        """
        其他初始化操作，可直接使用self.dto_class和self.api_suffix
        """

修改后的子类代码 #

"""
locust_specific_dto.py
"""
from locust_base import BaseUser
# 导入对应的DTO类
from some_module import SpecificDto

class SpecificDtoUser(BaseUser):
    # 直接覆写类属性
    dto_class = SpecificDto
    api_suffix = "specificEndpoint"

    # 无需重写__init__，直接继承父类逻辑
    # 可添加子类专属的task等逻辑

修改后运行locust -f locust_specific_dto.py即可正常执行，同时完美实现逻辑复用。

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