如何对用于培训的JupyterHub实例进行自动化负载测试？

Locust完全能满足你的需求——它不仅能处理HTTP请求模拟用户登录JupyterHub，还能通过扩展支持WebSocket通信（Jupyter Notebook执行单元格的核心机制就是WebSocket），完美复现真实学生的操作流程。

核心原因：为什么requests不够？ #

Jupyter Notebook的单元格加载、执行不是通过简单的HTTP POST/GET完成的，而是通过WebSocket通道与内核（Kernel）交互，传递execute_request等指令，同时接收执行结果。纯requests无法处理这种双向实时通信，而Locust可以结合websocket-client库实现这一逻辑。

环境准备 #

先安装依赖：

pip install locust websocket-client

Locust测试脚本示例 #

下面是一个完整的测试脚本，包含模拟用户登录JupyterHub、获取Notebook服务器连接、执行单元格的完整流程：

from locust import HttpUser, task, between
import websocket
import json
import time
from urllib.parse import urljoin

class JupyterHubUser(HttpUser):
    wait_time = between(1, 3)  # 模拟用户操作间隔
    notebook_path = "test_notebook.ipynb"  # 替换为你的测试Notebook路径
    username = "test_user"  # 替换为测试用户名
    password = "test_pass"  # 替换为测试密码

    def on_start(self):
        # 登录JupyterHub
        self.login()
        # 获取Notebook服务器地址
        self.notebook_server_url = self.get_notebook_server_url()
        # 获取运行中的Kernel ID
        self.kernel_id = self.get_kernel_id()

    def login(self):
        # 访问登录页获取CSRF Token
        login_page = self.client.get("/hub/login")
        csrf_token = login_page.cookies.get("_xsrf")
        
        # 提交登录请求
        login_data = {
            "username": self.username,
            "password": self.password,
            "_xsrf": csrf_token
        }
        response = self.client.post("/hub/login", data=login_data)
        # 验证登录成功（正常会重定向到用户主页）
        assert response.status_code in (200, 302), "登录失败"

    def get_notebook_server_url(self):
        # 拼接用户Notebook服务器的基础地址
        return urljoin(self.client.base_url, f"/user/{self.username}/")

    def get_kernel_id(self):
        # 请求Kernel列表接口，获取第一个运行中的Kernel ID
        kernels_url = urljoin(self.notebook_server_url, "api/kernels")
        response = self.client.get(kernels_url)
        kernels = response.json()
        assert len(kernels) > 0, "未找到运行的Kernel，请确保Notebook已启动"
        return kernels[0]["id"]

    @task
    def run_notebook_cell(self):
        # 建立WebSocket连接到Kernel通信通道
        ws_url = self.notebook_server_url.replace("http", "ws") + f"api/kernels/{self.kernel_id}/channels"
        ws = websocket.create_connection(ws_url)
        
        try:
            # 构造并发送单元格执行请求
            execute_msg = {
                "header": {
                    "msg_id": f"locust_{int(time.time())}",
                    "msg_type": "execute_request"
                },
                "content": {
                    "code": "print('Hello from Locust!'); x = 1 + 2; x",  # 要执行的测试代码
                    "silent": False,
                    "store_history": True
                },
                "metadata": {},
                "parent_header": {}
            }
            ws.send(json.dumps(execute_msg))
            
            # 监听执行结果，直到收到执行完成的回复
            while True:
                response = ws.recv()
                msg = json.loads(response)
                if msg["msg_type"] == "execute_reply":
                    assert msg["content"]["status"] == "ok", "单元格执行失败"
                    break
                time.sleep(0.1)
        finally:
            ws.close()

运行测试 #

1. 将脚本保存为jupyter_locust.py
2. 启动Locust服务：

locust -f jupyter_locust.py --host=https://your-jupyterhub-domain.com

3. 打开浏览器访问http://localhost:8089，设置并发用户数、孵化速率，启动测试即可。

其他可选工具 #

• k6：支持HTTP和WebSocket的性能测试工具，用JavaScript编写，适合非Python栈的用户。
• nbclient：可直接批量执行Notebook内容，配合Locust的并发调度，适合复杂Notebook场景的测试。

内容的提问来源于stack exchange，提问作者Steven R Brandt