以库形式运行Locust时出现未知命令行参数问题求助

方案1：先解析自定义参数并清理sys.argv #

这种方法会先单独处理你的自定义命令行参数，然后从sys.argv中移除这些参数，避免Locust的解析器识别到它们。

修改后的代码示例：

import argparse
import sys
import gevent
from locust import HttpUser, task
from locust.env import Environment
from locust.stats import stats_printer, stats_history
from locust.log import setup_logging

# 先解析自定义参数
custom_parser = argparse.ArgumentParser()
custom_parser.add_argument("-c", dest="custom_arg")
custom_parser.add_argument("-b", dest="another_arg")
# 分离已解析的自定义参数和剩余参数（如Locust相关参数）
args, remaining_args = custom_parser.parse_known_args()
# 更新sys.argv，只保留脚本名和剩余参数
sys.argv = [sys.argv[0]] + remaining_args

setup_logging("INFO", None)

class MyUser(HttpUser):
    host = "https://docs.locust.io"

    @task
    def t(self):
        self.client.get("/")

env = Environment(user_classes=[MyUser])
runner = env.create_local_runner()

web_ui = env.create_web_ui("127.0.0.1", 8089)

env.events.init.fire(environment=env, runner=runner, web_ui=web_ui)

gevent.spawn(stats_printer(env.stats))
gevent.spawn(stats_history, env.runner)

runner.start(1, spawn_rate=10)
gevent.spawn_later(60, lambda: runner.quit())

runner.greenlet.join()
web_ui.stop()

你可以通过args.custom_arg和args.another_arg获取自定义参数的值，同时Locust的解析器只会处理剩余的参数，不会再报错。

方案2：将自定义参数添加到Locust的解析器中 #

这种方法把你的自定义参数注册到Locust的内置命令行解析器里，让Locust能够识别这些参数，避免冲突。

修改后的代码示例：

import gevent
from locust import HttpUser, task
from locust.env import Environment
from locust.stats import stats_printer, stats_history
from locust.log import setup_logging
from locust import argument_parser

# 向Locust的解析器添加自定义参数
parser = argument_parser.get_parser()
parser.add_argument("-c", dest="custom_arg", help="自定义参数c")
parser.add_argument("-b", dest="another_arg", help="自定义参数b")

setup_logging("INFO", None)

class MyUser(HttpUser):
    host = "https://docs.locust.io"

    @task
    def t(self):
        self.client.get("/")

env = Environment(user_classes=[MyUser])
runner = env.create_local_runner()

web_ui = env.create_web_ui("127.0.0.1", 8089)

# 获取解析后的自定义参数值
parsed_args = argument_parser.parse_args()
print(f"自定义参数c的值：{parsed_args.custom_arg}")
print(f"自定义参数b的值：{parsed_args.another_arg}")

env.events.init.fire(environment=env, runner=runner, web_ui=web_ui)

gevent.spawn(stats_printer(env.stats))
gevent.spawn(stats_history, env.runner)

runner.start(1, spawn_rate=10)
gevent.spawn_later(60, lambda: runner.quit())

runner.greenlet.join()
web_ui.stop()

方案选择建议 #

• 如果你的程序已有复杂的自定义参数解析逻辑，建议使用方案1，保持参数解析的独立性。
• 如果你的自定义参数较少，且希望和Locust的参数体系整合，建议使用方案2。

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