Python

[Locust] User 클래스에 커스텀 arguments 전달

담다디다 2022. 4. 29. 10:15

locust를 실행할 때 user 클래스에게 특정 값을 초기에 전달해야하는 경우가 있다.

예를 들면, 특정 도메인으로 요청을 보내고 싶을 때가 있고 아닐 때가 있는데

이 분기를 실행할 때 user 클래스로 True/False를 전달해서 결정할 수 있도록 하는 것이다.

 

 

이번 포스팅은 locust 실행 시 커스텀 인수를 전달하는 2가지 방법에 대해 정리하고자 한다.

 

  1. command line 명령어를 커스터마이징 하기
  2. 환경 변수 설정하기

 

 

 

1. command line 명령어 커스터마이징 하기

 

 

locust에는 디폴트로 설정된 커맨드라인 명령어가 있다. 여기에 내가 커스터마이징 한 명령어를 추가하는 방법을 사용한다

 

이 방법은 locust의 events 데코레이션을 활용하게 된다.

 

 

from locust import events

@events.init_command_line_parser.add_listener	#데코레이터 추가
def set_command_line(parser):
    #명령어 추가
    parser.add_argument("--test", type=bool, default=False, help="let's test customizing")

 

 

코드를 한 줄씩 살펴보면 먼저 events 데코레이션에서 커맨드라인 parser의 add_listener를 사용한다.

 

그리고 parser를 파라미터로 전달받아 argument를 추가한다.

 

add_argument 내부 파라미터는 여러가지가 있는데, 차례로 명령어, 입력타입, 디폴트값, 그리고 설명을 위 코드에 넣었다.

이외의 유용한 파라미터들은 공식문서를 통해 파악하면 된다.

 

위 함수는 user클래스 밖에 선언하면 된다.

 

재밌는 점은, 위와 같이 커맨드라인 명령어를 추가하고 locust -h 를 누르면 디폴트 명령어들과 함께 내가 추가한 명령어까지 콘솔에 나온다는 것이다. help는 이때 출력되는 명령어 설명에 들어가는 부분이다.

 

 

 

위와 같이 명령어를 추가했다면 user 클래스에서 사용해보자.

 

 

class TestUser(HttpUser):
    if self.environment.parsed_options.test == True:	#명령어 입력값 가져오기
    	self.client.get("/test")

 

 

사용은 간단하게 environment에 포함된 parsed_options에서 내가 추가한 명령어의 값을 가져오면 된다.

필자는 --test 라고 작성했으니 사용도 parsed_options.test로 사용하면 되는 것이다.

 

 

 

이제 이 명령어를 콘솔에서 사용해보자

 

 

locust --test True

 

 

 

 

2.  환경 변수 설정하기

 

따로 명령어를 추가하고 사용하는 것이 번거롭다면 단순히 환경변수를 사용하는 방법도 있다.

 

환경변수의 장점은 어떤 코드에서든 입력값에 접근할 수 있다는 점이다.

필자는 TaskSet에서 입력값을 가져와야 했는데 추가한 명령어로 접근하는 경우 user class에서 값을 가져와 TaskSet으로 전달하는 번거로운 과정을 거쳐야 했다. TaskSet은 environment에 접근할 수 없기 때문이다.

 

그러나 환경변수를 사용하면 TaskSet이든 일반 함수든 자유롭게 입력값에 접근할 수 있다.

 

 

환경변수는 따로 추가하거나 설정해줄 필요 없이 콘솔에서 다음과 같이 실행하면 된다.

 

 

<linux/mac>
TEST=True locust ...

<Window>
SET TEST=True
locust ...

 

 

운영체제에 따라 달라지긴 하지만 위처럼 간단하게 환경변수를 추가할 수 있다.

 

이렇게 추가한 변수를 사용해보자

 

 

import os
from locust import TaskSet, task

class TestTaskSet(TaskSet):
	
    @task
    def test_task(self):
    	if os.environ['TEST'] == "True":		# 환경변수 값 가져오기
        	self.client.get("/test")

 

 

사용은 os에서 환경변수를 가져오는 방식과 동일하다.

단, 환경변수는 타입이 str이므로 필요하면 타입을 변경하는 방식으로 사용하면 된다.

반응형