Merge pull request #2497 from locustio/typing-things

Update mypy, some type hints, and some tiny updates to Runner classes

locust/env.py

@@ -30,10 +30,10 @@ def __init__(
         user_classes: Optional[List[Type[User]]] = None,
         shape_class: Optional[LoadTestShape] = None,
         tags: Optional[List[str]] = None,
-        locustfile: str = None,
+        locustfile: Optional[str] = None,
         exclude_tags: Optional[List[str]] = None,
-        events: Events = None,
-        host: str = None,
+        events: Optional[Events] = None,
+        host: Optional[str] = None,
         reset_stats=False,
         stop_timeout: Optional[float] = None,
         catch_exceptions=True,

locust/runners.py

@@ -316,6 +316,10 @@ def start(
     ) -> None:
         ...
 
+    @abstractmethod
+    def send_message(self, msg_type: str, data: Optional[Any] = None, client_id: Optional[str] = None) -> None:
+        ...
+
     def start_shape(self) -> None:
         """
         Start running a load test with a custom LoadTestShape specified in the :meth:`Environment.shape_class <locust.env.Environment.shape_class>` parameter.
@@ -431,15 +435,15 @@ class LocalRunner(Runner):
     Runner for running single process load test
     """
 
-    # always set to 0 for LocalRunner
-    worker_index = 0
-
     def __init__(self, environment) -> None:
         """
         :param environment: Environment instance
         """
         super().__init__(environment)
-
+        # These attributes dont make a lot of sense for LocalRunner
+        # but it makes it easier to write tests that work for both local and distributed runs
+        self.worker_index = 0
+        self.client_id = socket.gethostname() + "_" + uuid4().hex
         # Only when running in standalone mode (non-distributed)
         self._local_worker_node = WorkerNode(id="local")
         self._local_worker_node.user_classes_count = self.user_classes_count
@@ -451,7 +455,9 @@ def on_user_error(user_instance, exception, tb):
 
         self.environment.events.user_error.add_listener(on_user_error)
 
-    def _start(self, user_count: int, spawn_rate: float, wait: bool = False, user_classes: list = None) -> None:
+    def _start(
+        self, user_count: int, spawn_rate: float, wait: bool = False, user_classes: Optional[list] = None
+    ) -> None:
         """
         Start running a load test
 
@@ -553,7 +559,7 @@ def stop(self) -> None:
             return
         super().stop()
 
-    def send_message(self, msg_type: str, data: Optional[Any] = None) -> None:
+    def send_message(self, msg_type: str, data: Optional[Any] = None, client_id: Optional[str] = None) -> None:
         """
         Emulates internodal messaging by calling registered listeners
 
@@ -1357,12 +1363,15 @@ def stats_reporter(self) -> NoReturn:
                 logger.error(f"Temporary connection lost to master server: {e}, will retry later.")
             gevent.sleep(WORKER_REPORT_INTERVAL)
 
-    def send_message(self, msg_type: str, data: Optional[Dict[str, Any]] = None) -> None:
+    def send_message(
+        self, msg_type: str, data: Optional[Dict[str, Any]] = None, client_id: Optional[str] = None
+    ) -> None:
         """
         Sends a message to master node
 
         :param msg_type: The type of the message to send
         :param data: Optional data to send
+        :param client_id: (unused)
         """
         logger.debug("Sending %s message to master" % msg_type)
         self.client.send(Message(msg_type, data, self.client_id))

tox.ini

@@ -49,5 +49,5 @@ deps = flake8==6.0.0
 commands = flake8 . --count --show-source --statistics
 
 [testenv:mypy]
-deps = mypy==0.981
+deps = mypy==1.7.1
 commands = mypy locust/