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+
21+ # base_url = "http://135.237.2.132"
22+ # client_secret = "AL1igbD8V6c9YjrmkapVL0vdtwvI7aP4wAWbCoqYEMw="
23+ # client_key = "UID2-C-I-146-h/Ju95.+SM3BKYyEIkWcGPK2OOVXCzoGezunDCKgeXYs="
24+ email_count = 5000
25+
26+ class IdentityMapTasks (TaskSet ):
27+ @task
28+ def identity_map_large_batch_task (self ):
29+ try :
30+ response = self .user .identityMapClient .generate_identity_map (
31+ IdentityMapInput .from_emails ([f"test{ i } @example.com" for i in range (email_count )])
32+ )
33+ if response .is_success ():
34+ self .user .environment .runner .stats .log_request ("IdentityMap" , "generate_identity_map" , response .elapsed_time * 1000 , len (response .response )) # Log successful request
35+ else :
36+ self .user .environment .runner .stats .log_error ("IdentityMap" , "generate_identity_map" , f"Failed with status: { response .status } " )
37+ except Exception as e :
38+ self .user .environment .runner .stats .log_error ("IdentityMap" , "generate_identity_map" , f"Exception: { e } " )
39+ print (f"Error in identity_map_large_batch_task: { e } " )
40+
41+ class IdentityMapUser (User ):
42+ wait_time = between (1 , 2 )
43+
44+ def __init__ (self , * args , ** kwargs ):
45+ super ().__init__ (* args , ** kwargs )
46+ self .identityMapClient = IdentityMapClient (
47+ base_url ,
48+ client_key ,
49+ client_secret )
50+ self .tasks = [IdentityMapTasks ] # Assign TaskSet
51+
52+ # Handle summary data
53+ def on_locust_stop (environment , ** kwargs ):
54+ if isinstance (environment .runner , MasterRunner ):
55+ summary = {
56+ "start_time" : environment .runner .start_time ,
57+ "end_time" : time .time (),
58+ "duration" : time .time () - environment .runner .start_time ,
59+ "stats" : environment .runner .stats .serialize_stats (),
60+ }
61+ with open ("locust_summary.json" , "w" ) as f :
62+ json .dump (summary , f , indent = 4 )
63+ print ("Locust summary saved to locust_summary.json" )
64+
65+ events .quitting .add_listener (on_locust_stop )
0 commit comments