forked from yuvipanda/hubtraf
-
Notifications
You must be signed in to change notification settings - Fork 0
/
__main__.py
99 lines (91 loc) · 2.9 KB
/
__main__.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
import asyncio
import structlog
import argparse
import random
import time
import socket
from hubtraf.user import User, OperationError
from hubtraf.auth.dummy import login_dummy
from functools import partial
async def simulate_user(hub_url, username, password, delay_seconds, code_execute_seconds):
await asyncio.sleep(delay_seconds)
async with User(username, hub_url, partial(login_dummy, password=password)) as u:
try:
await u.login()
await u.ensure_server()
await u.start_kernel()
await u.assert_code_output("5 * 4", "20", 5, code_execute_seconds)
except OperationError:
pass
finally:
try:
if u.state == User.States.KERNEL_STARTED:
await u.stop_kernel()
except OperationError:
# We'll try to sto the server anyway
pass
try:
if u.state == User.States.SERVER_STARTED:
await u.stop_server()
except OperationError:
# Nothing to do
pass
def main():
argparser = argparse.ArgumentParser()
argparser.add_argument(
'hub_url',
help='Hub URL to send traffic to (without a trailing /)'
)
argparser.add_argument(
'user_count',
type=int,
help='Number of users to simulate'
)
argparser.add_argument(
'--user-prefix',
default=socket.gethostname(),
help='Prefix to use when generating user names'
)
argparser.add_argument(
'--user-session-min-runtime',
default=60,
type=int,
help='Min seconds user is active for'
)
argparser.add_argument(
'--user-session-max-runtime',
default=300,
type=int,
help='Max seconds user is active for'
)
argparser.add_argument(
'--user-session-max-start-delay',
default=60,
type=int,
help='Max seconds by which all users should have logged in'
)
argparser.add_argument(
'--json',
action='store_true',
help='True if output should be JSON formatted'
)
args = argparser.parse_args()
processors=[structlog.processors.TimeStamper(fmt="ISO")]
if args.json:
processors.append(structlog.processors.JSONRenderer())
else:
processors.append(structlog.dev.ConsoleRenderer())
structlog.configure(processors=processors)
awaits = []
for i in range(args.user_count):
awaits.append(simulate_user(
args.hub_url,
f'{args.user_prefix}-' + str(i),
'hello',
int(random.uniform(0, args.user_session_max_start_delay)),
int(random.uniform(args.user_session_min_runtime, args.user_session_max_runtime))
))
loop = asyncio.get_event_loop()
loop.run_until_complete(asyncio.gather(*awaits))
if __name__ == '__main__':
main()