added mypy to tox

fixed typing errors in locust/test.

.gitignore

@@ -3,11 +3,14 @@ locust.wpr
 locust.egg-info/**
 locustio.egg-info/**
 locust/_version.py
+locust/test/mock_*.py
 docs/_build/**
 docs/cli-help-output.txt
 docs/config-options.rst
 mock.*.egg
 web_test_*.csv
+err.txt
+out.txt
 .eggs/
 dist/**
 .idea/**

locust/test/__init__.py

@@ -3,6 +3,13 @@
 
     # work around occasional "zmq.error.ZMQError: Too many open files"
     # this is done in main.py when running locust proper so we need to do it here as well
-    resource.setrlimit(resource.RLIMIT_NOFILE, [10000, resource.RLIM_INFINITY])
+    resource.setrlimit(
+        resource.RLIMIT_NOFILE,
+        (
+            10000,
+            resource.RLIM_INFINITY,
+        ),
+    )
+    changed_rlimit = True
 except Exception:
-    pass  # Some OS:es will not allow changing NOFILE, but let's ignore that
+    changed_rlimit = False

locust/test/mock_logging.py

@@ -1,12 +1,17 @@
 import logging
 
+from typing import List, Union, Dict
+from types import TracebackType
+
+LogMessage = List[Union[str, Dict[str, TracebackType]]]
+
 
 class MockedLoggingHandler(logging.Handler):
-    debug = []
-    warning = []
-    info = []
-    error = []
-    critical = []
+    debug: List[LogMessage] = []
+    warning: List[LogMessage] = []
+    info: List[LogMessage] = []
+    error: List[LogMessage] = []
+    critical: List[LogMessage] = []
 
     def emit(self, record):
         if record.exc_info:

locust/test/test_dispatch.py

@@ -1,7 +1,7 @@
 import time
 import unittest
 from operator import attrgetter
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Type
 
 from locust import User
 from locust.dispatch import UsersDispatcher
@@ -3189,7 +3189,7 @@ def __str__(self):
                 self.fixed_counts, self.weights, self.target_user_count
             )
 
-    def case_handler(self, cases: List[RampUpCase], expected: Dict[str, int], user_classes: List[User]):
+    def case_handler(self, cases: List[RampUpCase], expected: List[Dict[str, int]], user_classes: List[Type[User]]):
         self.assertEqual(len(cases), len(expected))
 
         for case_num in range(len(cases)):
@@ -3448,7 +3448,7 @@ def _aggregate_dispatched_users(d: Dict[str, Dict[str, int]]) -> Dict[str, int]:
 
 
 def _user_count(d: Dict[str, Dict[str, int]]) -> int:
-    return sum(map(sum, map(dict.values, d.values())))
+    return sum(map(sum, map(dict.values, d.values())))  # type: ignore
 
 
 def _user_count_on_worker(d: Dict[str, Dict[str, int]], worker_node_id: str) -> int:

locust/test/test_log.py

@@ -10,6 +10,7 @@
 from locust.log import greenlet_exception_logger
 from .testcases import LocustTestCase
 from .util import temporary_file
+from . import changed_rlimit
 
 
 class TestGreenletExceptionLogger(LocustTestCase):
@@ -128,8 +129,10 @@ def my_task(self):
                 stderr=subprocess.STDOUT,
                 timeout=10,
             ).decode("utf-8")
-        # might give warning about not being able to change system open file limit (e.g. Windows10 + WSL2)
-        self.assertTrue(output.strip().endswith("running my_task"))
+        if not changed_rlimit:
+            self.assertTrue(output.strip().endswith("running my_task"))
+        else:
+            self.assertEqual("running my_task", output.strip())
 
     def test_log_to_file(self):
         with temporary_file(

locust/test/test_runners.py

@@ -45,7 +45,7 @@
     User,
     task,
 )
-from retry import retry
+from retry import retry  # type: ignore
 from .util import patch_env
 
 NETWORK_BROKEN = "network broken"
@@ -1386,7 +1386,7 @@ def tick(self):
 
             master.stop()
 
-    @unittest.skip
+    @unittest.skip(reason="takes a lot of time and has randomness to it")
     def test_distributed_shape_fuzzy_test(self):
         """
         Incredibility useful test to find issues with dispatch logic. This test allowed to find

locust/user/task.py

@@ -2,7 +2,7 @@
 import random
 import traceback
 from time import time
-from typing import TYPE_CHECKING, Callable, Any, List, Union, TypeVar, Optional, overload
+from typing import TYPE_CHECKING, Callable, List, Union, TypeVar, Optional, Type, overload
 from typing_extensions import final
 
 import gevent
@@ -15,7 +15,7 @@
 
 
 logger = logging.getLogger(__name__)
-TaskT = TypeVar("TaskT", bound=Callable[[Any], None])
+TaskT = TypeVar("TaskT", bound=Union[Callable[..., None], Type["TaskSet"]])
 
 LOCUST_STATE_RUNNING, LOCUST_STATE_WAITING, LOCUST_STATE_STOPPING = ["running", "waiting", "stopping"]
 
@@ -43,6 +43,16 @@ def read_thread(self):
             @task(7)
             def create_thread(self):
                 pass
+
+            @task(25)
+            class ForumThread(TaskSet):
+                @task
+                def get_author(self):
+                    pass
+
+                @task
+                def get_created(self):
+                    pass
     """
 
     def decorator_func(func):

tox.ini

@@ -17,6 +17,7 @@ deps =
     pyquery
     cryptography
     black==22.1.0
+    mypy
 allowlist_externals =
     bash
     timeout
@@ -25,6 +26,7 @@ commands =
     flake8 . --count --show-source --statistics
     coverage run -m unittest discover []
     black --check .
+    mypy --ignore-missing-imports locust/
     bash -ec 'PYTHONUNBUFFERED=1 timeout 2s python3 examples/debugging.py >out.txt 2>err.txt || true'
     grep -m 1 '/hello' out.txt
     bash -ec '! grep . err.txt' # should be empty