dcache/bench/locustfile.py

import json
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


class HelloGrpcUser(adaptor.GrpcUser):
    stub_class = DcacheServiceStub
    leader = "localhost:9001"
    host = leader
    captcha_id = "locust"
    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)

    #   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,
                },
            }
        }

        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)


#        msg = dcache.CaptchaID(id=captcha_id)
#        resp = self.stub.AddVisitor(msg)
#        pprint(resp)