diff --git a/bench/locustfile.py b/bench/locustfile.py index b12524d..c846dfc 100644 --- a/bench/locustfile.py +++ b/bench/locustfile.py @@ -1,4 +1,5 @@ import json +import grpc from pprint import pprint from locust import FastHttpUser, between, task @@ -103,17 +104,38 @@ import gevent import adaptor from locust import events, task +host = "localhost:9001" +captcha_id = "locust" + + +def add_captcha(stub: DcacheServiceStub, captcha_id: str): + msg = dcache.AddCaptchaRequest( + id=captcha_id, + mcaptcha=dcache.MCaptcha( + duration=30, + defense=dcache.Defense( + levels=[ + dcache.Level(visitor_threshold=50, difficulty_factor=500), + dcache.Level(visitor_threshold=5000, difficulty_factor=50000), + ] + ), + ), + ) + resp = stub.AddCaptcha(msg) + pprint(f"Captcha added {captcha_id}: {resp}") + + +with grpc.insecure_channel(host) as channel: + stub = DcacheServiceStub(channel) + add_captcha(stub=stub, captcha_id=captcha_id) + class HelloGrpcUser(adaptor.GrpcUser): stub_class = DcacheServiceStub - leader = "localhost:9001" - host = leader - captcha_id = "locust" + host = host + captcha_id = captcha_id msg = RaftRequest(data=json.dumps({"AddVisitor": captcha_id})) - def on_start(self): - self.add_captcha(captcha_id=self.captcha_id) - def add_vote(self, captcha_id: str): resp = self.stub.Write(self.msg) @@ -121,28 +143,6 @@ class HelloGrpcUser(adaptor.GrpcUser): # resp = self.pipeline_write(data=self.pipeline_vote) # pprint(resp) - def add_captcha(self, captcha_id: str): - params = { - "AddCaptcha": { - "id": captcha_id, - "mcaptcha": { - "visitor_threshold": 0, - "defense": { - "levels": [ - {"visitor_threshold": 50, "difficulty_factor": 500}, - {"visitor_threshold": 5000, "difficulty_factor": 50000}, - ], - "current_visitor_threshold": 0, - }, - "duration": 30, - }, - } - } - - msg = RaftRequest(data=json.dumps(params)) - resp = self.stub.Write(msg) - pprint(f"Captcha added {captcha_id}: {resp}") - @task def addVote(self): self.add_vote(self.captcha_id)