locus-load.py

import base64
import json
import random
from datetime import date, timedelta

import requests
from locust import HttpLocust, TaskSet, task

events = []


def status(l):
    payload = random.choice(events)
    headers = {'content-type': 'application/json', 'x-client': 'locus', 'X-Proxy': 'http://clt-guard-09:9290'}
    l.client.post("/merge", headers=headers, data=json.dumps(payload))


class UserBehavior(TaskSet):

    @task(1)
    def on_start(self):
        status(self)


class WebsiteUser(HttpLocust):
    task_set = UserBehavior
    min_wait = 5000
    max_wait = 9000

    def get_events(self, app_version, date_from, date_to, q_from, q_size):
        query = '{"query": {"filtered": {"query": {"bool": {"must": [{"match": {"producer.version": "' + app_version + '"}}, {"match": {"producer.name": "guard"}}, {"match": {"type": "INCOMING_HTTP_REQUEST"}}]}},"filter": {"range": {"timestamp": {"gte": "' + date_from + '","lte": "' + date_to + '"}}}}},"size": ' + str(
            q_size) + ',"from": ' + str(q_from) + '}'

        _response = requests.post("http://clt-elastic-01.despexds.net:9290/audit/event/_search", query)
        _res = []
        if _response.status_code == 200:
            data = json.loads(_response.text)
            elems = data["hits"]["hits"];
            for elem in elems:
                _res.append(json.loads(base64.b64decode(elem["_source"]["content"]["body"])))
            return _res
        else:
            return []

    def setup(self):
        print("setup: preparing data for load test")
        global events
        date_from = str(date.today() - timedelta(2))
        date_to = str(date.today())
        app_version = "1.15.4"
        events = self.get_events(app_version, date_from, date_to, 0, 10000)

        print("setup done! events: ", len(events))