|  | 
|  | 1 | +# | 
|  | 2 | +# SEE https://docs.locust.io/en/stable/quickstart.html | 
|  | 3 | +# | 
|  | 4 | +# This script allows testing running a function via the map endpoint | 
|  | 5 | +# | 
|  | 6 | + | 
|  | 7 | + | 
|  | 8 | +import json | 
|  | 9 | +import random | 
|  | 10 | +from datetime import timedelta | 
|  | 11 | +from typing import Final | 
|  | 12 | +from uuid import UUID | 
|  | 13 | + | 
|  | 14 | +import jsf | 
|  | 15 | +from common.base_user import OsparcWebUserBase | 
|  | 16 | +from locust import events, task | 
|  | 17 | +from locust.argument_parser import LocustArgumentParser | 
|  | 18 | +from tenacity import ( | 
|  | 19 | +    Retrying, | 
|  | 20 | +    retry_if_exception_type, | 
|  | 21 | +    stop_after_delay, | 
|  | 22 | +    wait_exponential, | 
|  | 23 | +) | 
|  | 24 | + | 
|  | 25 | +_MAX_NJOBS: Final[int] = 50 | 
|  | 26 | + | 
|  | 27 | + | 
|  | 28 | +# Register the custom argument with Locust's parser | 
|  | 29 | +@events.init_command_line_parser.add_listener | 
|  | 30 | +def _(parser: LocustArgumentParser) -> None: | 
|  | 31 | +    parser.add_argument( | 
|  | 32 | +        "--function-uuid", | 
|  | 33 | +        type=UUID, | 
|  | 34 | +        default=None, | 
|  | 35 | +        help="The function UUID to test", | 
|  | 36 | +    ) | 
|  | 37 | +    parser.add_argument( | 
|  | 38 | +        "--function-input-json-schema", | 
|  | 39 | +        type=str, | 
|  | 40 | +        default=None, | 
|  | 41 | +        help="JSON schema for the function job inputs", | 
|  | 42 | +    ) | 
|  | 43 | +    parser.add_argument( | 
|  | 44 | +        "--max-poll-time-seconds", | 
|  | 45 | +        type=int, | 
|  | 46 | +        default=60, | 
|  | 47 | +        help="Maximum time to wait for the function job collection to complete", | 
|  | 48 | +    ) | 
|  | 49 | +    parser.add_argument( | 
|  | 50 | +        "--n-jobs", | 
|  | 51 | +        type=int, | 
|  | 52 | +        default=None, | 
|  | 53 | +        help=f"Number of jobs to run via map-endpoint. If not set, a random number between 0 and {_MAX_NJOBS} is selected", | 
|  | 54 | +    ) | 
|  | 55 | + | 
|  | 56 | + | 
|  | 57 | +class WebApiUser(OsparcWebUserBase): | 
|  | 58 | +    @task | 
|  | 59 | +    def map_function(self) -> None: | 
|  | 60 | + | 
|  | 61 | +        function_uuid = self.environment.parsed_options.function_uuid | 
|  | 62 | +        if function_uuid is None: | 
|  | 63 | +            raise ValueError("function-uuid argument is required") | 
|  | 64 | +        if self.environment.parsed_options.function_input_json_schema is None: | 
|  | 65 | +            raise ValueError("function-input-json-schema argument is required") | 
|  | 66 | +        job_input_schema = json.loads( | 
|  | 67 | +            self.environment.parsed_options.function_input_json_schema | 
|  | 68 | +        ) | 
|  | 69 | +        max_poll_time = timedelta( | 
|  | 70 | +            seconds=self.environment.parsed_options.max_poll_time_seconds | 
|  | 71 | +        ) | 
|  | 72 | +        n_jobs = ( | 
|  | 73 | +            int(self.environment.parsed_options.n_jobs) | 
|  | 74 | +            if self.environment.parsed_options.n_jobs is not None | 
|  | 75 | +            else random.randint(1, _MAX_NJOBS) | 
|  | 76 | +        ) | 
|  | 77 | + | 
|  | 78 | +        # map function | 
|  | 79 | +        job_input_faker = jsf.JSF(job_input_schema) | 
|  | 80 | +        response = self.authenticated_post( | 
|  | 81 | +            url=f"/v0/functions/{function_uuid}:map", | 
|  | 82 | +            json=[job_input_faker.generate() for _ in range(n_jobs)], | 
|  | 83 | +            headers={ | 
|  | 84 | +                "x-simcore-parent-project-uuid": "null", | 
|  | 85 | +                "x-simcore-parent-node-id": "null", | 
|  | 86 | +            }, | 
|  | 87 | +            name="/v0/functions/[function_uuid]:map", | 
|  | 88 | +        ) | 
|  | 89 | +        response.raise_for_status() | 
|  | 90 | +        job_collection_uuid = response.json().get("uid") | 
|  | 91 | + | 
|  | 92 | +        # wait for the job to complete | 
|  | 93 | +        for attempt in Retrying( | 
|  | 94 | +            stop=stop_after_delay(max_delay=max_poll_time), | 
|  | 95 | +            wait=wait_exponential(multiplier=1, min=1, max=10), | 
|  | 96 | +            reraise=True, | 
|  | 97 | +            retry=retry_if_exception_type(ValueError), | 
|  | 98 | +        ): | 
|  | 99 | +            with attempt: | 
|  | 100 | +                job_status_response = self.authenticated_get( | 
|  | 101 | +                    f"/v0/function_job_collections/{job_collection_uuid}/status", | 
|  | 102 | +                    name="/v0/function_job_collections/[job_collection_uuid]/status", | 
|  | 103 | +                ) | 
|  | 104 | +                job_status_response.raise_for_status() | 
|  | 105 | +                all_job_statuses = job_status_response.json().get("status") | 
|  | 106 | +                assert isinstance(all_job_statuses, list) | 
|  | 107 | +                if any(status != "SUCCESS" for status in all_job_statuses): | 
|  | 108 | +                    raise ValueError( | 
|  | 109 | +                        f"Function job ({job_collection_uuid=}) for function ({function_uuid=}) returned {all_job_statuses=}" | 
|  | 110 | +                    ) | 
0 commit comments