ruff

Author: danopolan

src/grasshopper/lib/grasshopper.py

@@ -12,9 +12,10 @@
 if os.name == "posix":
     import resource  # pylint: disable=import-error
 
+from functools import wraps
+
 import gevent
 import locust
-from functools import wraps
 from grasshopper.lib.journeys.base_journey import BaseJourney
 from grasshopper.lib.util.listeners import GrasshopperListeners
 from locust import LoadTestShape
@@ -177,37 +178,47 @@ def launch_test(
         # processed because eventually, we will need to consult the shape to get the max
         # number of virtual users
         Grasshopper._assign_weights_to_user_classes(weighted_user_classes)
-        
+
         # Set up iteration limit if specified
         iterations = kwargs.get("iterations", 0)
         runtime = kwargs.get("runtime")
-        
+
         logger.info(f"Test run limits: {runtime} seconds and {iterations} iterations.")
         # Log test limits at the start
         if runtime and iterations > 0:
             logger.info("Test will stop when either limit is reached first.")
         elif runtime:
-            logger.info(f"Test will stop when runtime limit of {runtime} seconds will be reached.")
+            logger.info(
+                f"Test will stop when runtime limit of {runtime} seconds will be reached."
+            )
         elif iterations > 0:
-            logger.info(f"Test will stop when ietration limit of {iterations} will be reached.")
-        
+            logger.info(
+                f"Test will stop when ietration limit of {iterations} will be reached."
+            )
+
         if iterations > 0:
             Grasshopper._setup_iteration_limit(env, iterations)
-        
+
         env.runner.start_shape()
-        
+
         # Always set runtime-based termination if runtime is specified
         # This allows both runtime and iterations to work together (whichever happens first)
         if runtime:
+
             def handle_runtime_limit():
                 # Check if iteration limit was already reached
-                if not (hasattr(env.runner, 'iteration_target_reached') and env.runner.iteration_target_reached):
+                if not (
+                    hasattr(env.runner, "iteration_target_reached")
+                    and env.runner.iteration_target_reached
+                ):
                     # Runtime limit reached first
-                    logger.info(f"Test stopped: Runtime limit of {runtime} seconds reached.")
+                    logger.info(
+                        f"Test stopped: Runtime limit of {runtime} seconds reached."
+                    )
                 os.kill(os.getpid(), signal.SIGINT)
-            
+
             gevent.spawn_later(runtime, handle_runtime_limit)
-        
+
         env.runner.greenlet.join()
 
         return env
@@ -220,38 +231,44 @@ def _assign_weights_to_user_classes(weighted_user_classes):
     @staticmethod
     def _setup_iteration_limit(env: Environment, iterations: int):
         """Set up iteration limiting for the test.
-        
+
         This method patches TaskSet.execute_task to track iterations and stop
         the test when the iteration limit is reached.
-        
+
         Args:
             env: The Locust Environment object
             iterations: Maximum number of iterations to run
         """
         runner = env.runner
         runner.iterations_started = 0
         runner.iteration_target_reached = False
-        
+
         def iteration_limit_wrapper(method):
             @wraps(method)
             def wrapped(self, task):
                 if runner.iterations_started >= iterations:
                     if not runner.iteration_target_reached:
                         runner.iteration_target_reached = True
-                        logger.info(f"Test stopped: Iteration limit of {iterations} reached.")
+                        logger.info(
+                            f"Test stopped: Iteration limit of {iterations} reached."
+                        )
                     if runner.user_count == 1:
                         logger.info("Last user stopped, quitting runner")
                         # Send final stats and quit
-                        gevent.spawn_later(0.1, lambda: os.kill(os.getpid(), signal.SIGINT))
+                        gevent.spawn_later(
+                            0.1, lambda: os.kill(os.getpid(), signal.SIGINT)
+                        )
                     raise StopUser()
                 runner.iterations_started = runner.iterations_started + 1
                 method(self, task)
-            
+
             return wrapped
-        
+
         # Patch TaskSet methods to add iteration limiting
         TaskSet.execute_task = iteration_limit_wrapper(TaskSet.execute_task)
-        DefaultTaskSet.execute_task = iteration_limit_wrapper(DefaultTaskSet.execute_task)
+        DefaultTaskSet.execute_task = iteration_limit_wrapper(
+            DefaultTaskSet.execute_task
+        )
 
     @staticmethod
     def load_shape(shape_name: str, **kwargs) -> LoadTestShape:

tests/unit/test_iterations.py

@@ -1,22 +1,23 @@
 """Unit tests for the iterations feature."""
 
+from unittest.mock import MagicMock, create_autospec, patch
+
 import pytest
-from unittest.mock import MagicMock, patch, create_autospec
 from grasshopper.lib.grasshopper import Grasshopper
 from grasshopper.lib.journeys.base_journey import BaseJourney
 from locust.env import Environment
+from locust.exception import StopUser
 from locust.runners import LocalRunner
 from locust.user.task import DefaultTaskSet, TaskSet
-from locust.exception import StopUser
 
 
 class MockJourney(BaseJourney):
     """Mock journey for testing."""
-    
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.task_count = 0
-    
+
     def example_task(self):
         self.task_count += 1
 
@@ -35,24 +36,24 @@ def mock_environment():
 def test_setup_iteration_limit_initializes_runner_attributes(mock_environment):
     """Test that _setup_iteration_limit properly initializes runner attributes."""
     Grasshopper._setup_iteration_limit(mock_environment, 10)
-    
+
     assert mock_environment.runner.iterations_started == 0
-    assert mock_environment.runner.iteration_target_reached == False
+    assert mock_environment.runner.iteration_target_reached is False
 
 
 def test_setup_iteration_limit_patches_taskset_methods(mock_environment):
     """Test that _setup_iteration_limit patches TaskSet methods."""
     # Store original methods
     original_taskset_execute = TaskSet.execute_task
     original_default_taskset_execute = DefaultTaskSet.execute_task
-    
+
     try:
         Grasshopper._setup_iteration_limit(mock_environment, 10)
-        
+
         # Check that methods have been patched (wrapped)
         assert TaskSet.execute_task != original_taskset_execute
         assert DefaultTaskSet.execute_task != original_default_taskset_execute
-        
+
     finally:
         # Restore original methods
         TaskSet.execute_task = original_taskset_execute
@@ -63,104 +64,98 @@ def test_iteration_limit_stops_after_reaching_limit(mock_environment):
     """Test that tasks stop executing after iteration limit is reached."""
     # Store original method
     original_taskset_execute = TaskSet.execute_task
-    
+
     try:
         Grasshopper._setup_iteration_limit(mock_environment, 2)
-        
+
         # Create a mock TaskSet
         mock_taskset = MagicMock(spec=TaskSet)
         mock_task = MagicMock()
-        
+
         # Execute tasks up to the limit
         TaskSet.execute_task(mock_taskset, mock_task)
         assert mock_environment.runner.iterations_started == 1
-        
+
         TaskSet.execute_task(mock_taskset, mock_task)
         assert mock_environment.runner.iterations_started == 2
-        
+
         # Next execution should raise StopUser
         with pytest.raises(StopUser):
             TaskSet.execute_task(mock_taskset, mock_task)
-        
+
         # Verify that iteration target was reached
-        assert mock_environment.runner.iteration_target_reached == True
-        
+        assert mock_environment.runner.iteration_target_reached is True
+
     finally:
         # Restore original method
         TaskSet.execute_task = original_taskset_execute
 
 
 def test_launch_test_with_iterations_no_runtime_timer():
     """Test that launch_test doesn't set runtime timer when iterations is specified."""
-    with patch.object(Grasshopper, 'set_ulimit'), \
-         patch.object(Grasshopper, '_setup_iteration_limit') as mock_setup_iterations, \
-         patch('grasshopper.lib.grasshopper.Environment') as MockEnvironment, \
-         patch('grasshopper.lib.grasshopper.gevent.spawn_later') as mock_spawn_later, \
-         patch('grasshopper.lib.grasshopper.gevent.spawn'), \
-         patch('grasshopper.lib.grasshopper.locust.stats.stats_history'), \
-         patch('grasshopper.lib.grasshopper.GrasshopperListeners'):
-        
+    with patch.object(Grasshopper, "set_ulimit"), patch.object(
+        Grasshopper, "_setup_iteration_limit"
+    ) as mock_setup_iterations, patch(
+        "grasshopper.lib.grasshopper.Environment"
+    ) as MockEnvironment, patch(
+        "grasshopper.lib.grasshopper.gevent.spawn_later"
+    ) as mock_spawn_later, patch("grasshopper.lib.grasshopper.gevent.spawn"), patch(
+        "grasshopper.lib.grasshopper.locust.stats.stats_history"
+    ), patch("grasshopper.lib.grasshopper.GrasshopperListeners"):
         # Setup mock environment
         mock_env = MagicMock()
         mock_env.runner = MagicMock()
         mock_env.runner.greenlet = MagicMock()
         mock_env.runner.greenlet.join = MagicMock()
         MockEnvironment.return_value = mock_env
-        
+
         # Mock shape instance
         mock_shape = MagicMock()
         mock_shape.configured_runtime = 120
-        
+
         # Launch test with iterations
-        kwargs = {
-            'iterations': 10,
-            'runtime': 120,
-            'shape_instance': mock_shape
-        }
-        
+        kwargs = {"iterations": 10, "runtime": 120, "shape_instance": mock_shape}
+
         Grasshopper.launch_test(MockJourney, **kwargs)
-        
+
         # Verify that _setup_iteration_limit was called
         mock_setup_iterations.assert_called_once_with(mock_env, 10)
-        
+
         # Verify that runtime timer was NOT set (spawn_later not called)
         mock_spawn_later.assert_not_called()
 
 
 def test_launch_test_without_iterations_sets_runtime_timer():
     """Test that launch_test sets runtime timer when iterations is not specified."""
-    with patch.object(Grasshopper, 'set_ulimit'), \
-         patch.object(Grasshopper, '_setup_iteration_limit') as mock_setup_iterations, \
-         patch('grasshopper.lib.grasshopper.Environment') as MockEnvironment, \
-         patch('grasshopper.lib.grasshopper.gevent.spawn_later') as mock_spawn_later, \
-         patch('grasshopper.lib.grasshopper.gevent.spawn'), \
-         patch('grasshopper.lib.grasshopper.locust.stats.stats_history'), \
-         patch('grasshopper.lib.grasshopper.GrasshopperListeners'):
-        
+    with patch.object(Grasshopper, "set_ulimit"), patch.object(
+        Grasshopper, "_setup_iteration_limit"
+    ) as mock_setup_iterations, patch(
+        "grasshopper.lib.grasshopper.Environment"
+    ) as MockEnvironment, patch(
+        "grasshopper.lib.grasshopper.gevent.spawn_later"
+    ) as mock_spawn_later, patch("grasshopper.lib.grasshopper.gevent.spawn"), patch(
+        "grasshopper.lib.grasshopper.locust.stats.stats_history"
+    ), patch("grasshopper.lib.grasshopper.GrasshopperListeners"):
         # Setup mock environment
         mock_env = MagicMock()
         mock_env.runner = MagicMock()
         mock_env.runner.greenlet = MagicMock()
         mock_env.runner.greenlet.join = MagicMock()
         MockEnvironment.return_value = mock_env
-        
+
         # Mock shape instance
         mock_shape = MagicMock()
         mock_shape.configured_runtime = 120
-        
+
         # Launch test without iterations (or with 0)
-        kwargs = {
-            'iterations': 0,
-            'runtime': 120,
-            'shape_instance': mock_shape
-        }
-        
+        kwargs = {"iterations": 0, "runtime": 120, "shape_instance": mock_shape}
+
         Grasshopper.launch_test(MockJourney, **kwargs)
-        
+
         # Verify that _setup_iteration_limit was NOT called
         mock_setup_iterations.assert_not_called()
-        
+
         # Verify that runtime timer WAS set
         mock_spawn_later.assert_called_once()
         # Check that it was called with the runtime value
-        assert mock_spawn_later.call_args[0][0] == 120
+        assert mock_spawn_later.call_args[0][0] == 120