Context manager / Decorator for adding manual Locust stats entries

locust_manual_report.py

"""
Usage examples:

# As context manager:
with manual_report("stats_entry_name"):
    # do stuff

# As a decorator for tasks:
@task
@manual_report
def my_task(self):
    time.sleep(random.random())    
"""


from contextlib import contextmanager
from time import time

from locust import events


@contextmanager
def _manual_report(name):
    start_time = time()
    try:
        yield
    except Exception as e:
        events.request_failure.fire(
            request_type="manual",
            name=name,
            response_time=(time()-start_time) * 1000,
            response_length=0,
            exception=e,
        )
        raise
    else:
        events.request_success.fire(
            request_type="manual",
            name=name,
            response_time=(time()-start_time) * 1000,
            response_length=0,
        ) 

def manual_report(name_or_func):
    if callable(name_or_func):
        # used as decorator without name argument specified
        return _manual_report(name_or_func.__name__)(name_or_func)
    else:
        return _manual_report(name_or_func)