I currently have this locust test:

import logging
from locust import HttpUser, TaskSet, task, constant

log = logging.getLogger("rest-api-performance-test")


def get_headers():
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
    }
    return headers


def headers_token(auth_token):
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        "auth-token": str(auth_token),
    }
    return headers


class LoginTasks(TaskSet):
    def post_login(self):
        headers = get_headers()
        login_response = self.client.post("/api/auth",
                                          json={"key": "[email protected]", "secret": "cstrong"},
                                          headers=headers, catch_response=False, name="Login")
        login_data = login_response.json()
        auth_token_value = login_data["results"][0]["auth-token"]
        return auth_token_value


class UserTasks(UserTasks):
    @task
    def test_get_list_of_workers(self):
        auth_token = self.post_login()
        try:
            with self.client.get("/api/client/workers",
                                 headers=headers_token(auth_token), catch_response=True,
                                 name="Get Worker Lists") as request_response:
                if request_response.status_code == 200:
                    assert (
                            '"label": "Salena Mead S"' in request_response.text
                    )
                    request_response.success()
                    log.info("API call resulted in success.")

                else:
                    request_response.failure(request_response.text)
                    log.error("API call resulted in failed.")

        except Exception as e:
            log.error(f"Exception occurred! details are {e}")


class WebsiteUser(HttpUser):
    host = "https://test.com"
    wait_time = constant(1)
    tasks = [UserTasks]

the tests runs as expected but post_login is required by multiple tests since is the one who generates the authentication token used by most of the APIs that I'm testing, is there a way to avoid use inheritance from class LoginTasks and find a better solution? The reason I want to avoid it is post_login is not the only method that is going be used many times so I don't want to use multiple inheritance on my UserTasks class.

Any help is appreciated.

CodePudding user response：

Move the function out of the class and pass in the client you want it to use.

def post_login(client):
    headers = get_headers()
    login_response = client.post("/api/auth",
    … 

You can then call it when you need it the same way you call get_headers().

auth_token = post_login(self.client)