1+ import json
2+ import time
3+ import sys
4+ from locust import User , task , between , events , TaskSet
5+ from locust .runners import MasterRunner
6+ from uid2_client import IdentityMapClient , IdentityMapInput
7+
8+ def _usage ():
9+ print ('Usage: venv/bin/python -m locust -f performance-testing/uid2-operator/locust-identity-map.py <base_url> <client_key> <client_secret>'
10+ , file = sys .stderr )
11+ sys .exit (1 )
12+
13+
14+ if len (sys .argv ) <= 4 :
15+ _usage ()
16+
17+ base_url = sys .argv [1 ]
18+ client_key = sys .argv [2 ]
19+ client_secret = sys .argv [3 ]
20+ email_count = 5000
21+
22+ class IdentityMapTasks (TaskSet ):
23+ @task
24+ def identity_map_large_batch_task (self ):
25+ try :
26+ response = self .user .identityMapClient .generate_identity_map (
27+ IdentityMapInput .from_emails ([f"test{ i } @example.com" for i in range (email_count )])
28+ )
29+ if response .is_success ():
30+ self .user .environment .runner .stats .log_request ("IdentityMap" , "generate_identity_map" , response .elapsed_time * 1000 , len (response .response )) # Log successful request
31+ else :
32+ self .user .environment .runner .stats .log_error ("IdentityMap" , "generate_identity_map" , f"Failed with status: { response .status } " )
33+ except Exception as e :
34+ self .user .environment .runner .stats .log_error ("IdentityMap" , "generate_identity_map" , f"Exception: { e } " )
35+ print (f"Error in identity_map_large_batch_task: { e } " )
36+
37+ class IdentityMapUser (User ):
38+ wait_time = between (1 , 2 )
39+
40+ def __init__ (self , * args , ** kwargs ):
41+ super ().__init__ (* args , ** kwargs )
42+ self .identityMapClient = IdentityMapClient (
43+ base_url ,
44+ client_key ,
45+ client_secret )
46+ self .tasks = [IdentityMapTasks ] # Assign TaskSet
47+
48+ # Handle summary data
49+ def on_locust_stop (environment , ** kwargs ):
50+ if isinstance (environment .runner , MasterRunner ):
51+ summary = {
52+ "start_time" : environment .runner .start_time ,
53+ "end_time" : time .time (),
54+ "duration" : time .time () - environment .runner .start_time ,
55+ "stats" : environment .runner .stats .serialize_stats (),
56+ }
57+ with open ("locust_summary.json" , "w" ) as f :
58+ json .dump (summary , f , indent = 4 )
59+ print ("Locust summary saved to locust_summary.json" )
60+
61+ events .quitting .add_listener (on_locust_stop )
0 commit comments