Merge pull request #3301 from locustio/simplify-custom-messages-match-case-statements-in-runners.py

refactor case statements and update to use 3.10 set syntax

Author: cyberw

locust/opentelemetry.py

@@ -15,10 +15,8 @@ def setup_opentelemetry() -> bool:
         logger.error("OpenTelemetry SDK is not installed, opentelemetry not enabled. Run 'pip install locust[otel]'")
         return False
 
-    traces_exporters = set(e.strip().lower() for e in os.getenv("OTEL_TRACES_EXPORTER", "otlp").split(",") if e.strip())
-    metrics_exporters = set(
-        e.strip().lower() for e in os.getenv("OTEL_METRICS_EXPORTER", "otlp").split(",") if e.strip()
-    )
+    traces_exporters = {e.strip().lower() for e in os.getenv("OTEL_TRACES_EXPORTER", "otlp").split(",") if e.strip()}
+    metrics_exporters = {e.strip().lower() for e in os.getenv("OTEL_METRICS_EXPORTER", "otlp").split(",") if e.strip()}
 
     if traces_exporters == {"none"} and metrics_exporters == {"none"}:
         logger.info("No OpenTelemetry exporters configured, opentelemetry not enabled")

locust/runners.py

@@ -1149,23 +1149,22 @@ def handle_message(self, client_id: str, msg: Message) -> None:
                             self.quit()
             case "exception":
                 self.log_exception(msg.node_id, msg.data["msg"], msg.data["traceback"])
+            case _ if lc := self.custom_messages.get(msg.type):
+                listener, concurrent = lc
+                logger.debug(
+                    f"Received {msg.type} message from worker {msg.node_id} (index {self.get_worker_index(msg.node_id)})"
+                )
+                try:
+                    if not concurrent:
+                        listener(environment=self.environment, msg=msg)
+                    else:
+                        gevent.spawn(listener, environment=self.environment, msg=msg)
+                except Exception:
+                    logging.error(f"Uncaught exception in handler for {msg.type}\n{traceback.format_exc()}")
             case _:
-                if lc := self.custom_messages.get(msg.type):
-                    listener, concurrent = lc
-                    logger.debug(
-                        f"Received {msg.type} message from worker {msg.node_id} (index {self.get_worker_index(msg.node_id)})"
-                    )
-                    try:
-                        if not concurrent:
-                            listener(environment=self.environment, msg=msg)
-                        else:
-                            gevent.spawn(listener, environment=self.environment, msg=msg)
-                    except Exception:
-                        logging.error(f"Uncaught exception in handler for {msg.type}\n{traceback.format_exc()}")
-                else:
-                    logger.warning(
-                        f"Unknown message type received from worker {msg.node_id} (index {self.get_worker_index(msg.node_id)}): {msg.type}"
-                    )
+                logger.warning(
+                    f"Unknown message type received from worker {msg.node_id} (index {self.get_worker_index(msg.node_id)}): {msg.type}"
+                )
 
         self.check_stopped()
 
@@ -1423,16 +1422,15 @@ def handle_message(self, msg: Message) -> None:
             case "spawning_complete":
                 # master says we have finished spawning (happens only once during a normal rampup)
                 self.environment.events.spawning_complete.fire(user_count=msg.data["user_count"])
-            case _:
-                if lc := self.custom_messages.get(msg.type):
-                    listener, concurrent = lc
-                    logger.debug(f"Received {msg.type} message from master")
-                    if not concurrent:
-                        listener(environment=self.environment, msg=msg)
-                    else:
-                        gevent.spawn(listener, self.environment, msg)
+            case _ if lc := self.custom_messages.get(msg.type):
+                listener, concurrent = lc
+                logger.debug(f"Received {msg.type} message from master")
+                if not concurrent:
+                    listener(environment=self.environment, msg=msg)
                 else:
-                    logger.warning(f"Unknown message type received: {msg.type}")
+                    gevent.spawn(listener, self.environment, msg)
+            case _:
+                logger.warning(f"Unknown message type received: {msg.type}")
 
     def stats_reporter(self) -> NoReturn:
         while True:

locust/test/test_runners.py

@@ -3366,7 +3366,11 @@ def my_task(self):
             sleep(0.1)
             server.mocked_send(Message("stats", BAD_MESSAGE, "zeh_fake_client1"))
             messages = server.get_messages()
-            self.assertEqual(5, len(messages))
+            self.assertEqual(messages[0].type, "ack")
+            self.assertEqual(messages[1].type, "spawn")
+            self.assertEqual(messages[2].type, "spawning_complete")
+            self.assertEqual(messages[3].type, "reconnect")
+            self.assertEqual(messages[4].type, "ack")
 
             # Expected message order in outbox: ack, spawn, reconnect, ack
             self.assertEqual(

locust/user/markov_taskset.py

@@ -206,7 +206,7 @@ def dfs(task_name):
                 dfs(dest)
 
     dfs(tasks[0].__name__)
-    unreachable = set([task.__name__ for task in tasks]) - visited
+    unreachable = {task.__name__ for task in tasks} - visited
 
     if len(unreachable) > 0:
         logging.warning(f"The following markov tasks are unreachable in class {classname}: {unreachable}")

locust/web.py

@@ -257,51 +257,52 @@ def swarm() -> Response:
             user_count = None
             spawn_rate = None
             for key, value in request.form.items():
-                if key == "user_count":  # if we just renamed this field to "users" we wouldn't need this
-                    user_count = int(value)
-                    parsed_options_dict["users"] = user_count
-                elif key == "spawn_rate":
-                    spawn_rate = float(value)
-                    parsed_options_dict[key] = spawn_rate
-                elif key == "host":
-                    # Replace < > to guard against XSS
-                    environment.host = str(request.form["host"]).replace("<", "").replace(">", "")
-                    parsed_options_dict[key] = environment.host
-                elif key == "user_classes":
-                    # Set environment.parsed_options.user_classes to the selected user_classes
-                    parsed_options_dict[key] = request.form.getlist("user_classes")
-                elif key == "run_time":
-                    if not value:
-                        continue
-                    try:
-                        run_time = parse_timespan(value)
-                        parsed_options_dict[key] = run_time
-                    except ValueError:
-                        err_msg = "Valid run_time formats are : 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc."
-                        logger.error(err_msg)
-                        return jsonify({"success": False, "message": err_msg, "host": environment.host})
-                elif key == "profile":
-                    environment.profile = str(request.form["profile"]) or None
-                    parsed_options_dict[key] = environment.profile
-                elif key in parsed_options_dict:
-                    # update the value in environment.parsed_options, but dont change the type.
-                    parsed_options_value = parsed_options_dict[key]
-
-                    if isinstance(parsed_options_value, bool):
-                        parsed_options_dict[key] = value == "true"
-                    elif parsed_options_value is None:
-                        parsed_options_dict[key] = value
-                    elif isinstance(parsed_options_value, list):
-                        if "," in value:
-                            value_as_list = value.split(",")
+                match key:
+                    case "user_count":  # if we just renamed this field to "users" we wouldn't need this
+                        user_count = int(value)
+                        parsed_options_dict["users"] = user_count
+                    case "spawn_rate":
+                        spawn_rate = float(value)
+                        parsed_options_dict[key] = spawn_rate
+                    case "host":
+                        # Replace < > to guard against XSS
+                        environment.host = str(request.form["host"]).replace("<", "").replace(">", "")
+                        parsed_options_dict[key] = environment.host
+                    case "user_classes":
+                        # Set environment.parsed_options.user_classes to the selected user_classes
+                        parsed_options_dict[key] = request.form.getlist("user_classes")
+                    case "run_time":
+                        if not value:
+                            continue
+                        try:
+                            run_time = parse_timespan(value)
+                            parsed_options_dict[key] = run_time
+                        except ValueError:
+                            err_msg = "Valid run_time formats are : 20, 20s, 3m, 2h, 1h20m, 3h30m10s, etc."
+                            logger.error(err_msg)
+                            return jsonify({"success": False, "message": err_msg, "host": environment.host})
+                    case "profile":
+                        environment.profile = str(request.form["profile"]) or None
+                        parsed_options_dict[key] = environment.profile
+                    case _ if key in parsed_options_dict:
+                        # update the value in environment.parsed_options, but dont change the type.
+                        parsed_options_value = parsed_options_dict[key]
+
+                        if isinstance(parsed_options_value, bool):
+                            parsed_options_dict[key] = value == "true"
+                        elif parsed_options_value is None:
+                            parsed_options_dict[key] = value
+                        elif isinstance(parsed_options_value, list):
+                            if "," in value:
+                                value_as_list = value.split(",")
+                            else:
+                                value_as_list = request.form.getlist(key)
+                            if all(isinstance(x, int) for x in parsed_options_value):
+                                parsed_options_dict[key] = list(map(int, value_as_list))
+                            else:
+                                parsed_options_dict[key] = value_as_list
                         else:
-                            value_as_list = request.form.getlist(key)
-                        if all(isinstance(x, int) for x in parsed_options_value):
-                            parsed_options_dict[key] = list(map(int, value_as_list))
-                        else:
-                            parsed_options_dict[key] = value_as_list
-                    else:
-                        parsed_options_dict[key] = type(parsed_options_value)(value)
+                            parsed_options_dict[key] = type(parsed_options_value)(value)
 
             if environment.shape_class and environment.runner is not None:
                 environment.runner.start_shape()