locust 2.19.2.dev35__py3-none-any.whl → 2.19.2.dev47__py3-none-any.whl

locust/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.19.2.dev35'
-__version_tuple__ = version_tuple = (2, 19, 2, 'dev35')
+__version__ = version = '2.19.2.dev47'
+__version_tuple__ = version_tuple = (2, 19, 2, 'dev47')

locust/rpc/protocol.py

@@ -1,4 +1,30 @@
 import msgpack
+import datetime
+from typing import Type
+
+try:
+    from bson import ObjectId  # type: ignore
+except ImportError:
+
+    class ObjectId:  # type: ignore
+        def __init__(self, s):
+            raise Exception("You need to install pymongo or at least bson to be able to send/receive ObjectIds")
+
+
+def decode(obj):
+    if "__datetime__" in obj:
+        obj = datetime.datetime.strptime(obj["as_str"], "%Y%m%dT%H:%M:%S.%f")
+    elif "__ObjectId__" in obj:
+        obj = ObjectId(obj["as_str"])
+    return obj
+
+
+def encode(obj):
+    if isinstance(obj, datetime.datetime):
+        return {"__datetime__": True, "as_str": obj.strftime("%Y%m%dT%H:%M:%S.%f")}
+    elif isinstance(obj, ObjectId):
+        return {"__ObjectId__": True, "as_str": str(obj)}
+    return obj
 
 
 class Message:
@@ -11,9 +37,9 @@ class Message:
         return f"<Message {self.type}:{self.node_id}>"
 
     def serialize(self):
-        return msgpack.dumps((self.type, self.data, self.node_id))
+        return msgpack.dumps((self.type, self.data, self.node_id), default=encode)
 
     @classmethod
     def unserialize(cls, data):
-        msg = cls(*msgpack.loads(data, raw=False, strict_map_key=False))
+        msg = cls(*msgpack.loads(data, raw=False, strict_map_key=False, object_hook=decode))
         return msg

locust/runners.py

@@ -37,6 +37,7 @@ import greenlet
 import psutil
 from gevent.event import Event
 from gevent.pool import Group
+import inspect
 
 from . import User
 from locust import __version__
@@ -380,7 +381,11 @@ class Runner:
         """
         if self.state == STATE_STOPPED:
             return
-        logger.debug("Stopping all users")
+        try:
+            caller = inspect.getframeinfo(inspect.stack()[1][0])
+            logger.debug(f"Stopping all users (called from {caller.filename}:{caller.lineno})")
+        except Exception:
+            logger.debug("Stopping all users (couldnt determine where stop() was called from)")
         self.environment.events.test_stopping.fire(environment=self.environment)
         self.final_user_classes_count = {**self.user_classes_count}
         self.update_state(STATE_CLEANUP)
@@ -1104,7 +1109,11 @@ class MasterRunner(DistributedRunner):
                 logger.debug(
                     f"Received {msg.type} message from worker {msg.node_id} (index {self.get_worker_index(msg.node_id)})"
                 )
-                self.custom_messages[msg.type](environment=self.environment, msg=msg)
+                try:
+                    self.custom_messages[msg.type](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}"

locust/web.py

@@ -227,8 +227,14 @@ class WebUI:
                         return jsonify({"success": False, "message": err_msg, "host": environment.host})
                 elif key in parsed_options_dict:
                     # update the value in environment.parsed_options, but dont change the type.
-                    # This won't work for parameters that are None
-                    parsed_options_dict[key] = type(parsed_options_dict[key])(value)
+                    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] = parsed_options_value
+                    else:
+                        parsed_options_dict[key] = type(parsed_options_dict[key])(value)
 
             if environment.shape_class and environment.runner is not None:
                 environment.runner.start_shape()