Python Load Testing

locust.py

from locust import HttpLocust, TaskSet, task
from locust.events import request_failure
import random, json, string

BASE_URL = "/api/v1"
USER_NAME = "username"
PASSWORD = "passwordhere"
#todo move to class var
headers = {
  "Accept": "*/*",
  "Connection": "keep-alive",
  "Pragma": "no-cache"
}

class UserBehavior(TaskSet):
  def random_string(self, length):
    return ''.join(random.choice(string.ascii_letters) for m in range(length))

  #todo: make this more robust
  def get_viable_user_id(self):
    return int(headers["X-User-Id"])-1

  def on_start(self):
    """ on_start is called when a Locust start before any task is scheduled """
    self.post_session()

  def on_stop(self):
    """ on_stop is called when the TaskSet is stopping """
    print("done")

  def on_failure(request_type, name, response_time, exception, **kwargs):
    print(exception.request.url)
    print(exception.response.status_code)
    print(exception.response.content)

  def post_session(self):
    post_data = {"username":USER_NAME, "email":USER_NAME, "password":PASSWORD}
    with self.client.post(f"{BASE_URL}/sessions", post_data, catch_response=True) as response:
      if "X-User-Id" not in response.headers:
        response.failure("login failed")

      global headers
      headers["X-User-Id"] = response.headers["X-User-Id"]
      headers["X-User-Authentication-Token"] = response.headers["X-User-Authentication-Token"]

  @task(1)
  def get_session_history(self):
    self.client.get(f"{BASE_URL}/session_history", headers=headers)

  @task(2)
  def get_favorites(self):
    self.client.get(f"{BASE_URL}/favorites?offset=0", headers=headers)

  @task(2)
  def post_check_favorite_limit(self):
    self.client.post(f"{BASE_URL}/favorites/check_favorite_limit?user_id=" + headers["X-User-Id"], headers=headers)

  @task(20)
  def get_matches(self):
    self.client.get(f"{BASE_URL}/matches?offset=0&count=144", headers=headers)

  @task(3)
  def get_matches_preview(self):
    self.client.get(f"{BASE_URL}/matches/preview", headers=headers)

  @task(30)
  def get_syncs(self):
    self.client.get(f"{BASE_URL}/syncs", headers=headers)

  @task(20)
  def get_own_user(self):
    self.client.get(f"{BASE_URL}/users/" + headers["X-User-Id"], headers=headers)

  @task(2)
  def post_user(self):
    post_data = {
      "user":{
        "ask_me":self.random_string(100), 
        "save_source":"profileStep3"
      }
    }
    self.client.patch(f"{BASE_URL}/users/" + headers["X-User-Id"], json=post_data, headers=headers)

  @task(3)
  def get_conversations(self):
    self.client.get(f"{BASE_URL}/conversations",  headers=headers)
    #get conversations with message id
    #http://localhost:3000/api/v1/conversations?message_id=723071

    #read

  @task(2)
  def get_plans(self):
    self.client.get(f"{BASE_URL}/plans", headers=headers)

  @task(30)
  def post_page_views(self):
    pages = [
      "/settings", 
      "/members", 
      "/conversations", 
      "/plans", 
      "/favorites/received",
      "/profilesetup/step/1",
      "/profilesetup/step/2",
      "/profilesetup/step/3",
      "/profilesetup/step/4",
      "/onboarding/tour"
    ]

    post_data = {
      "page_view": {
        "has_viewed_all": "false", 
        "path": random.choice(pages)
      }
    }
    self.client.post(f"{BASE_URL}/page_views", json=post_data, headers=headers)

  @task(2)
  def get_photos(self):
    self.client.get(f"{BASE_URL}/photos", headers=headers)

  @task(1)
  def post_regenerate_match_deck(self):
    self.client.post(f"{BASE_URL}/users/" + headers["X-User-Id"] + "/regenerate_match_deck", headers=headers)

  @task(5)
  def post_hidings(self):
    post_data = {
    "hiding":{
        "expected":"true",
        "diff":"0.000",
        "like":False,
        "daily_profile":random.randint(1,9),
        "source":"profile-light",
        "hidden_user_id":self.get_viable_user_id()
      }
    }
    with self.client.post(f"{BASE_URL}/hidings", json=post_data, headers=headers, catch_response=True) as response:
      if response.status_code == 400:
        print(response.content)
        if response.content and response.content.error == "User not found":
          response.success()
      elif response.status_code == 500 and 'sex_enum' in response.content.decode("utf-8"):
        response.success() #known issue

  @task(5)
  def post_favorites(self):
    self.client.post(f"{BASE_URL}/favorites?expected=false&diff=0.00&like=true&daily_profile=5&source=profile-light&user_id=" + str(self.get_viable_user_id()), headers=headers, name="/favorites")


class WebsiteUser(HttpLocust):
  task_set = UserBehavior
  min_wait = 1000
  max_wait = 15000