Locust: how I share auth cookie with the rest of tasks only for current locust user?

Question

I am looking into to moving my multi-threaded python script to locust.

A simple explanation of what my script does is:

Create a thread per user
In each thread authenticates user and get auth cookie
With that auth cookie perform various api calls at a set interval

When i started looking into locust, I have noticed that the only way to perform each task at its own specific interval, I would need to create a taskset per task.

This brought up an issue of how do i share the auth cookie for the given spawned user between task sets? Since in the long run I also need to share response data between taskset for the given spawned user as it differs between spawned users.

In the sample code below, all of the users spawned by locust, share the same "storage.cookie". Is there a way to keep storage.cookie unique per user, share it with all tasks sets for the given spawned user by locust ? Does locust report on which user is currently executing the task?

from __future__ import print_function
from locust import Locust, TaskSet, task, HttpLocust
import json


def auth(l):
    payload = {"username":"some_username","password":"some_password"} 
    resp = l.client.post('/auth', data = json.dumps(payload))
    storage.cookie = # get auth cookie from resp

def do_i_auth(l):
    if len(storage.cookie) == 0:
        auth(l)

class storage(object):
    cookie == ''

class first_call(TaskSet):
    def on_start(self):
        do_i_auth(self)

    @task
    def get_api_a(self):
        headers = {"Cookie":storage.cookie}
        self.client.get('/api_a', headers)

class second_call(TaskSet):
    def on_start(self):
        do_i_auth(self)

    @task
    def get_api_b(self):
        headers = {"Cookie":storage.cookie}
        self.client.get('/api_b', headers)

class api_A(HttpLocust):
    task_set = first_call
    min_wait = 5000
    max_wait = 5000    

class api_B(HttpLocust):
    task_set = second_call
    min_wait = 10000
    max_wait = 10000

score 6 · Answer 1 · answered Feb 21 '18 at 12:06

6

You can try have your authorization function return the cookie and have it stored in each class separately. Something like this:

from __future__ import print_function
from locust import Locust, TaskSet, task, HttpLocust
import json


def auth(l):
    payload = {"username":"some_username","password":"some_password"} 
    resp = l.client.post('/auth', data = json.dumps(payload))
    cookie = # get auth cookie from resp
    return cookie

class first_call(TaskSet):
    cookie = ""

    def on_start(self):
        self.cookie = auth(self)

    @task
    def get_api_a(self):
        headers = {"Cookie":self.cookie}
        self.client.get('/api_a', headers)

answered Feb 21 '18 at 12:06

Dimitris

204
2
5

1

there is a race condition here. the on_start runs in parallel with the task, you are just lucky if this works. – basickarl May 12 '21 at 16:14
No tasks will be started before on_start has finished, so there is no need to worry about a race condition. https://github.com/locustio/locust/blob/69f2a2b6eae03089562ff820ce86adc4b9107bff/locust/user/users.py#L134 – Cyberwiz Oct 13 '21 at 12:52

score 1 · Answer 2 · answered Feb 06 '19 at 18:48

I think the solution here is to not have separate classes for each call, but instead to have the calls as methods on a single class. That way you can store the cookies on the objects (referenced via self.cookie).

This worked for me:

https://gist.github.com/MatrixManAtYrService/1d83abd54adc9d4181f9ebb98b9799f7

score 1 · Answer 3 · edited Sep 03 '20 at 07:39

I recently implemented cookies in a DotNet application load test script.

Cookies should be passed using dictionary objects.

  cookiedict={}
  cookiedict['Key1'] = 'Value1'
  cookiedict['Key2'] = 'Value2'
  #Auth API
  self.auth_response = self.gettoken(cookiedict)  
  self.token = self.auth_response.cookies['RequestVerificationToken']
  self.cookies = self.auth_response.cookies
  #Login API
  cookiedict['RequestVerificationToken'] = self.token

` 
self.login_response=self.login(self.user_name,self.password,self.token,cookiedict)

Also note that you need to use HttpSession as well

from locust.clients import HttpSession
self.requests = HttpSession(consumer_cfg.rest_api_url)

 executor = self.requests.post
 if method == 'PUT':
    executor = self.requests.put
 elif method == 'GET':
    executor = self.requests.get


self._request_proceed(method='GET', url=url,  data=formdata,catch_response=catch_response, cookies = CookiesSent,allow_redirects = True)

Locust: how I share auth cookie with the rest of tasks only for current locust user?

3 Answers3