应用 Locust 快速上手写压测

引子

做为一个压测工具（库），locust 其实解决这么一个问题：AB 之类压测工具不能编写复杂的因果逻辑，而现实场景中，待压的服务往往是有一套完整执行流程的，比如 APP 要访问一个 API，是需要先鉴权（验明不是非 APP 访问），再登录换 Token，然后才是 API 调用……

这一切，在 locust 中都很容易实现，本质上，应用 locust 做压测，就是在写 Python 程序，只是它集成了一套不错的 UI，外加并发的benchmark功能。

至于写个压测为什么要用Python，是因为：这玩意心智负担低，你谷歌SO复制粘贴一把梭，直接上手就能写，大脑无需切换context，调试成本也低，没有比这语言更棒的了 🙂

骨架

一个 locust 压测程序的最简配置就是一个 .py 程序文件，按照约定，我们给它命名为 locustfile.py，它大概由这么两部分组成：

一个 TaskSet 类：里面放置你需要测试的各种任务，这些任务用 @task 装饰器来标记（洋气）
一个 Locust 类：它的一个实例就代表一个了用户，同时它提供了用户并发访问的能力；它的 task_set 成员需要引用 TaskSet 类，那是它要执行的任务。

当然，locustfile.py 还有很多实用的配置，如父子任务/等待时长/权重等设置，细节可参见官网文档。

举个栗子

我现在有一个 API 需要压测，它的服务端大概是这么设计的：

这个 API 地址是 http://example.org/api/dummy/test
任何访问这个 API 的请求，都需要在 Headers 的 Authorization 中提供登录后才会有的 Token
而要进行登录，需要提供这些 Headers：
- Content-Type: application/json
- timestamp：发起请求时APP 的时间戳
- nonce：一个随机数
- checksum： HMACSHA1(timestamp + nonce, secretKey)，是的，一次哈希，并且这个 secretKey 是仅 APP 和服务端才知道的密钥
- 以及正确的用户名/密码对
登录地址是 http://example.org/api/users/login，在验明正身后，会下发 JSON 格式的 JWT token

OK，因为够复杂，所以要写程序来做压测了，Python 干这个，工程复杂度正好

一个实现

环境：

Python：2.7.11
Locust： 0.7.3

代码：

from __future__ import print_function
from locust import HttpLocust, TaskSet, task
from hashlib import sha1
import time
import uuid
import hmac
import base64
import httplib
import json

secret = "__HOLY_SECRET_KEY__"
username = "Luke"
password = "I'mYourFather"
host = "http://example.org"
loginUrl = "/api/users/login"
apiUrl = "/api/dummy/test"

def hmacsha1(plaintext, secret):
    return hmac.new(secret, plaintext, sha1).hexdigest()
def timestamp():
    return str(int(time.time()))
def nonce():
    return str(uuid.uuid1())
def auth_header():
    ts = timestamp()
    nc = nonce()
    checksum = hmacsha1(ts + nc, secret)
    return {"Content-type": "application/json",
            "timestamp": ts,
             "nonce": nc,
             "checksum": checksum}

class MyTaskSet(TaskSet):
    def on_start(self):
        self.tokenInfo = None
        headers = auth_header()
        data = json.dumps({"username": username,"password": password})
        response = self.client.request(method="POST", url= loginUrl,
                                       data = data,
                                       headers = headers)
        print("LOGIN RESULT:", response.status_code, response.content)
        if response.status_code == 200:
            self.tokenInfo = json.loads(response.content)

    @task
    def dummy_test(self):
        if self.tokenInfo is not None:
            token = self.tokenInfo["token"]
            headers = auth_header()
            headers["Authorization"] = "Bearer " + token
            response = self.client.request(method="GET", url= apiUrl,
                                           headers = headers)
            print("API RESULT:", response.status_code, response.content)

class MyLocust(HttpLocust):
    task_set = MyTaskSet
    host = host

from __future__ import print_function

from locust import HttpLocust, TaskSet, task

from hashlib import sha1

import time

import uuid

import hmac

import base64

import httplib

import json

secret = "__HOLY_SECRET_KEY__"

username = "Luke"

password = "I'mYourFather"

host = "http://example.org"

loginUrl = "/api/users/login"

apiUrl = "/api/dummy/test"

def hmacsha1(plaintext, secret):

return hmac.new(secret, plaintext, sha1).hexdigest()

def timestamp():

return str(int(time.time()))

def nonce():

return str(uuid.uuid1())

def auth_header():

ts = timestamp()

nc = nonce()

checksum = hmacsha1(ts + nc, secret)

return {"Content-type": "application/json",

"timestamp": ts,

"nonce": nc,

"checksum": checksum}

class MyTaskSet(TaskSet):

def on_start(self):

self.tokenInfo = None

headers = auth_header()

data = json.dumps({"username": username,"password": password})

response = self.client.request(method="POST", url= loginUrl,

data = data,

headers = headers)

print("LOGIN RESULT:", response.status_code, response.content)

if response.status_code == 200:

self.tokenInfo = json.loads(response.content)

@task

def dummy_test(self):

if self.tokenInfo is not None:

token = self.tokenInfo["token"]

headers = auth_header()

headers["Authorization"] = "Bearer " + token

response = self.client.request(method="GET", url= apiUrl,

headers = headers)

print("API RESULT:", response.status_code, response.content)

class MyLocust(HttpLocust):

task_set = MyTaskSet

host = host

跑一个

先安装 Locust：

pip install locustio

1	pip install locustio

shell 里直接运行，locust 默认会运行 locustfile.py：

locust

locust

命令行会在服务起来后提示：

Starting web monitor at *:8089

1	Starting web monitor at *:8089

然后就可以在浏览器中访问 http://127.0.0.1:8089/：

locust-look-and-feel

输入希望同时访问的用户数，及每用户的秒并发量，开始压测：

locust-benchmark-new

在这个过程中，可以随时停止测试，调整参数，当然，也可以在测试数据收集完成后，导出测试结果：

locust-result-export

打赏作者

您的支持将激励我继续创作！

去打赏

Mo Ye

A Conservative

Published

June 24, 2017

Mo Ye in Engineering | June 24, 2017