from locust import constant, task, events
from locust.user.users import HttpUser


@events.init_command_line_parser.add_listener
def add_custom_arguments(parser):
    # add arguments that we should see in GUI
    parser.add_argument(
        "--arg1",
        type=str,
        help="Some argument 1",
        default="Foo",
    )

    parser.add_argument(
        "--arg2",
        type=int,
        help=f"Some argument 2 (default: {10})",
        default=10,
    )


class MyUser1(HttpUser):
    host = "http://example.com"
    wait_time = constant(1)

    @task
    def get_ip(self):
        print(
            "pasred custom args: "
            f"arg1={self.environment.parsed_options.arg1}, "
            f"arg2={self.environment.parsed_options.arg2}"
        )