153 lines
4.6 KiB
Python
153 lines
4.6 KiB
Python
import json
|
|
import grpc
|
|
from pprint import pprint
|
|
from locust import FastHttpUser, between, task
|
|
|
|
from dcache_py import dcache_pb2 as dcache
|
|
from dcache_py.dcache_pb2 import RaftRequest
|
|
from dcache_py.dcache_pb2_grpc import DcacheServiceStub
|
|
|
|
# password = "fooobarasdfasdf"
|
|
# username = "realaravinth"
|
|
#
|
|
#
|
|
# class Unprotected(FastHttpUser):
|
|
# # wait_time = between(5, 15)
|
|
# peers = [
|
|
# "http://localhost:9001",
|
|
# "http://localhost:9002",
|
|
# "http://localhost:9003",
|
|
# ]
|
|
# leader = "http://localhost:9001"
|
|
# host = leader
|
|
# captcha_id="locust"
|
|
#
|
|
# pipeline_vote = []
|
|
# for _ in range(0,1000):
|
|
# pipeline_vote.append({"AddVisitor": captcha_id})
|
|
#
|
|
## def on_start(self):
|
|
## resp = self.client.get(f"{self.leader}/metrics")
|
|
## data = resp.json()
|
|
## leader = data["Ok"]["membership_config"]["log_id"]["leader_id"]["node_id"]
|
|
## self.leader = self.peers[leader - 1]
|
|
## self.host = self.leader
|
|
## print(f"Leader: {self.host}")
|
|
## self.add_captcha(captcha_id="locust")
|
|
#
|
|
#
|
|
#
|
|
# def write(self, data):
|
|
# resp = self.client.post(f"{self.host}/write", json=data)
|
|
# print(f"RPC Status: {resp.status_code}")
|
|
# resp = resp.json()
|
|
# if "Err" in resp:
|
|
# leader = resp["Err"]["APIError"]["ForwardToLeader"]["leader_node"]["addr"]
|
|
# print(f"Forwarding write to leader {leader}")
|
|
# return write(leader, data)
|
|
# return resp["Ok"]["data"]
|
|
#
|
|
# def pipeline_write(self, data):
|
|
# resp = self.client.post(f"{self.host}/pipeline/write", json=data)
|
|
## print(f"RPC Status: {resp.status_code}")
|
|
# resp = resp.json()
|
|
# if "Err" in resp:
|
|
# leader = resp["Err"]["APIError"]["ForwardToLeader"]["leader_node"]["addr"]
|
|
# print(f"Forwarding write to leader {leader}")
|
|
# return write(leader, data)
|
|
# return resp
|
|
#
|
|
#
|
|
#
|
|
#
|
|
# def add_vote(self, captcha_id: str):
|
|
# resp = self.write(data={"AddVisitor": captcha_id})
|
|
# pprint(resp)
|
|
#
|
|
# def add_vote_pipeline(self, captcha_id: str):
|
|
# 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,
|
|
# },
|
|
# }
|
|
# }
|
|
# resp = self.write(data=params)
|
|
# pprint(f"Captcha added {captcha_id}: {resp}")
|
|
#
|
|
#
|
|
# @task
|
|
# def unprotected(self):
|
|
# self.add_vote_pipeline(captcha_id=self.captcha_id)
|
|
# ##self.add_vote(captcha_id="locust")
|
|
## data = {
|
|
## "username": username,
|
|
## "password": username,
|
|
## "confirm_password": username,
|
|
## }
|
|
## self.client.post("/unprotected", data=data)
|
|
#
|
|
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
|
|
host = host
|
|
captcha_id = captcha_id
|
|
msg = RaftRequest(data=json.dumps({"AddVisitor": captcha_id}))
|
|
|
|
def add_vote(self, captcha_id: str):
|
|
resp = self.stub.Write(self.msg)
|
|
|
|
# def add_vote_pipeline(self, captcha_id: str):
|
|
# resp = self.pipeline_write(data=self.pipeline_vote)
|
|
# pprint(resp)
|
|
|
|
@task
|
|
def addVote(self):
|
|
self.add_vote(self.captcha_id)
|
|
|
|
|
|
# msg = dcache.CaptchaID(id=captcha_id)
|
|
# resp = self.stub.AddVisitor(msg)
|
|
# pprint(resp)
|