Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
request_type="tavern",
name=name,
response_time=total_time,
exception=e
)
else:
total_time = int((time.time() - start_time) * 1000)
events.request_success.fire(
request_type="tavern",
name=name,
response_time=total_time,
response_length=0,
)
class TavernLocust(Locust):
def __init__(self, *args, **kwargs):
super(TavernLocust, self).__init__(*args, **kwargs)
self.client = TavernClient(self.host)
name = name[len(strip_prefix) :]
# Prefix relative URLs with the root path for readability.
if not name.startswith("http"):
name = "/" + name
# Replace query parameters with non-varying placeholders.
parsed = urllib.parse.urlparse(name)
query = urllib.parse.parse_qs(parsed.query)
for k in list(query.keys()):
query[k] = "X"
name = parsed._replace(query=urlencode(query)).geturl()
return name
class _AuthenticatedLocust(Locust):
"""Base for authenticated RDR REST API callers."""
def __init__(self, *args, **kwargs):
super(_AuthenticatedLocust, self).__init__(*args, **kwargs)
creds_file = os.environ["LOCUST_CREDS_FILE"]
instance = os.environ["LOCUST_TARGET_INSTANCE"]
# The "client" field gets copied to TaskSet instances.
self.client = _ReportingClient(creds_file=creds_file, default_instance=instance, parse_cli=False)
self.participant_generator = FakeParticipantGenerator(self.client, use_local_files=True)
class VersionCheckUser(_AuthenticatedLocust):
# 2 out of 100 users (Locust count of 100 recommended in load_test.sh).
weight = 2
# Hit the root/version endpoint every 10s.
min_wait = 1000 * 10 * weight
name = name[len(strip_prefix):]
# Prefix relative URLs with the root path for readability.
if not name.startswith('http'):
name = '/' + name
# Replace query parameters with non-varying placeholders.
parsed = urlparse.urlparse(name)
query = urlparse.parse_qs(parsed.query)
for k in query.keys():
query[k] = 'X'
name = parsed._replace(query=urlencode(query)).geturl()
return name
class _AuthenticatedLocust(Locust):
"""Base for authenticated RDR REST API callers."""
def __init__(self, *args, **kwargs):
super(_AuthenticatedLocust, self).__init__(*args, **kwargs)
creds_file = os.environ['LOCUST_CREDS_FILE']
instance = os.environ['LOCUST_TARGET_INSTANCE']
# The "client" field gets copied to TaskSet instances.
self.client = _ReportingClient(
creds_file=creds_file, default_instance=instance, parse_cli=False)
self.participant_generator = FakeParticipantGenerator(self.client, use_local_files=True)
class VersionCheckUser(_AuthenticatedLocust):
# 2 out of 100 users (Locust count of 100 recommended in load_test.sh).
weight = 2
# Hit the root/version endpoint every 10s.
min_wait = 1000 * 10 * weight
def get_access_token(self):
self.client.call_endpoint(
'GetAccessToken',
dict(owner=USER_FULL, scope=SCOPE),
self.session_token
)
@seq_task(4)
def cancel_session_token(self):
self.client.call_endpoint(
'CancelSessionToken',
dict(session_token=self.session_token),
)
class BrokerUser(Locust):
host = BROKER_HOST
task_set = JobSimulation
min_wait = 0
max_wait = 1
def __init__(self, *args, **kwargs):
self.client = BrokerClient(self.host, get_certificate())
name="Send_payload_"+task,
response_time=total_time,
exception=e)
raise e
else:
total_time = int((time.time() - start_time) * 1000)
events.request_success.fire(request_type="send_payload",
name="Send_payload_"+task,
response_time=total_time,
response_length=len(payload))
def disconnect(self):
self.ws.close()
class WSLocust(Locust):
def __init__(self, *args, **kwargs):
super(Locust, self)
self.client = WSLocustClient()
WORK_DIR = os.path.dirname(__file__)
# read kafka brokers from config
KAFKA_BROKERS = os.getenv("KAFKA_BROKERS", "kafka:9092").split(sep=",")
# read other environment variables
QUIET_MODE = True if os.getenv("QUIET_MODE", "true").lower() in ['1', 'true', 'yes'] else False
TASK_DELAY = int(os.getenv("TASK_DELAY", "0"))
# register additional logging handlers
if not QUIET_MODE:
events.request_success += additional_success_handler
events.request_failure += additional_failure_handler
class KafkaLocust(Locust):
client = None
def __init__(self, *args, **kwargs):
super(KafkaLocust, self).__init__(*args, **kwargs)
if not KafkaLocust.client:
KafkaLocust.client = KafkaClient(KAFKA_BROKERS)
class KafkaBehaviour(TaskSet):
def random_message(self, min_length=32, max_length=128):
return ''.join(random.choice(string.ascii_uppercase) for _ in range(random.randrange(min_length, max_length)))
def timestamped_message(self):
return f"{time.time() * 1000}:" + ("kafka" * 24)[:random.randint(32, 128)]
response_length = len(response[0]) + len(response[1])
events.request_success.fire(request_type=name, name=cmd, response_time=total_time, response_length=response_length)
else:
events.request_failure.fire(request_type=name, name=cmd, response_time=total_time, exception=e)
raise e
return response
return wrapper
class CustomUNIXSocketTransport(UNIXSocketTransport):
run_cmd = capture(UNIXSocketTransport.run_cmd)
class ErrandBoyLocust(Locust):
def __init__(self, *args, **kwargs):
super(ErrandBoyLocust, self).__init__(*args, **kwargs)
self.client = CustomUNIXSocketTransport()
class LoadTestTask(TaskSet):
@task
def run_cmd_ls(self):
self.client.run_cmd('ls -al')
@task
def run_cmd_date(self):
self.client.run_cmd('date')
class LoadTestUser(ErrandBoyLocust):
from locust import Locust, TaskSet, task
class MyTaskSet(TaskSet):
@task(20)
def hello(self):
pass
class Dummy(Locust):
task_set = MyTaskSet
min_wait = 5000
max_wait = 15000
start_time = time.time()
try:
result = func(*args, **kwargs)
except xmlrpclib.Fault as e:
total_time = int((time.time() - start_time) * 1000)
events.request_failure.fire(request_type="xmlrpc", name=name, response_time=total_time, exception=e)
else:
total_time = int((time.time() - start_time) * 1000)
events.request_success.fire(request_type="xmlrpc", name=name, response_time=total_time, response_length=0)
# In this example, I've hardcoded response_length=0. If we would want the response length to be
# reported correctly in the statistics, we would probably need to hook in at a lower level
return wrapper
class XmlRpcLocust(Locust):
"""
This is the abstract Locust class which should be subclassed. It provides an XML-RPC client
that can be used to make XML-RPC requests that will be tracked in Locust's statistics.
"""
def __init__(self, *args, **kwargs):
super(XmlRpcLocust, self).__init__(*args, **kwargs)
self.client = XmlRpcClient(self.host)
class ApiUser(XmlRpcLocust):
host = "http://127.0.0.1:8877/"
wait_time = between(0.1, 1)
class task_set(TaskSet):
@task(10)