-
Notifications
You must be signed in to change notification settings - Fork 6
/
3_load_testing_with_locust.py
151 lines (122 loc) · 4.43 KB
/
3_load_testing_with_locust.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
# %% [markdown]
"""
# Web API: 3. Load testing with Locust
This tutorial shows how to use an API endpoint
created in the FastAPI tutorial in load testing.
"""
# %pip install dff locust
# %% [markdown]
"""
## Running Locust
1. Run this file directly:
```bash
python {file_name}
```
2. Run locust targeting this file:
```bash
locust -f {file_name}
```
3. Run from python:
```python
import sys
from locust import main
sys.argv = ["locust", "-f", {file_name}]
main.main()
```
You should see the result at http://127.0.0.1:8089.
Make sure that your POST endpoint is also running (run the FastAPI tutorial).
"""
# %%
################################################################################
# this patch is only needed to run this file in IPython kernel
# and can be safely removed
import gevent.monkey
gevent.monkey.patch_all()
################################################################################
# %%
import uuid
import time
import sys
from locust import FastHttpUser, task, constant, main
from dff.script import Message
from dff.utils.testing import HAPPY_PATH, is_interactive_mode
# %%
class DFFUser(FastHttpUser):
wait_time = constant(1)
def check_happy_path(self, happy_path):
"""
Check a happy path.
For each `(request, response)` pair in `happy_path`:
1. Send request to the API endpoint and catch its response.
2. Compare API response with the `response`.
If they do not match, fail the request.
:param happy_path:
An iterable of tuples of
`(Message, Message | Callable(Message->str|None) | None)`.
If the second element is `Message`,
check that API response matches it.
If the second element is `None`,
do not check the API response.
If the second element is a `Callable`,
call it with the API response as its argument.
If the function returns a string,
that string is considered an error message.
If the function returns `None`,
the API response is considered correct.
"""
user_id = str(uuid.uuid4())
for request, response in happy_path:
with self.client.post(
f"/chat?user_id={user_id}",
headers={
"accept": "application/json",
"Content-Type": "application/json",
},
# Name is the displayed name of the request.
name=f"/chat?user_message={request.json()}",
data=request.json(),
catch_response=True,
) as candidate_response:
text_response = Message.model_validate(
candidate_response.json().get("response")
)
if response is not None:
if callable(response):
error_message = response(text_response)
if error_message is not None:
candidate_response.failure(error_message)
elif text_response != response:
candidate_response.failure(
f"Expected: {response.model_dump_json()}\n"
f"Got: {text_response.model_dump_json()}"
)
time.sleep(self.wait_time())
@task(3) # <- this task is 3 times more likely than the other
def dialog_1(self):
self.check_happy_path(HAPPY_PATH)
@task
def dialog_2(self):
def check_first_message(msg: Message) -> str | None:
if msg.text is None:
return f"Message does not contain text: {msg.model_dump_json()}"
if "Hi" not in msg.text:
return (
f'"Hi" is not in the response message: '
f"{msg.model_dump_json()}"
)
return None
self.check_happy_path(
[
# a function can be used to check the return message
(Message("Hi"), check_first_message),
# a None is used if return message should not be checked
(Message("i'm fine, how are you?"), None),
# this should fail
(Message("Hi"), check_first_message),
]
)
# %%
if __name__ == "__main__":
if is_interactive_mode():
sys.argv = ["locust", "-f", __file__]
main.main()