Simplify tests

Author: ajcasagrande

src/aiperf/timing/fixed_schedule_strategy.py

@@ -40,11 +40,11 @@ def __init__(
         # NOTE: This all needs to be set before the super call, because the base class will call
         # _setup_profiling_phase_config() which uses it to set the total expected requests.
 
-        # Reconstruct the full schedule from first_turn_timestamp and turn_delays
+        # Reconstruct the full schedule from the first turn timestamps
         self._schedule: list[tuple[int | float, str]] = []
         for conversation in dataset_metadata.conversations:
             if conversation.turns[0].timestamp_ms is not None:
-                # Add first turn
+                # Add first turn only, as the credit is for the whole conversation
                 self._schedule.append(
                     (conversation.turns[0].timestamp_ms, conversation.conversation_id)
                 )

src/aiperf/timing/timing_manager.py

@@ -89,7 +89,11 @@ async def _on_dataset_configured_notification(
         self, message: DatasetConfiguredNotification
     ) -> None:
         """Handle the dataset configured notification."""
-        self.debug(f"Received dataset configured notification: {message}")
+        self.debug(
+            lambda: f"Received dataset configured notification: {len(message.metadata.conversations)} conversations, "
+            f"{message.metadata.sampling_strategy.value} sampling strategy"
+        )
+
         self._dataset_metadata = message.metadata
         self._dataset_configured_event.set()
 

tests/unit/timing/test_benchmark_duration.py

@@ -66,6 +66,17 @@ def mixed_config(
     )
 
 
+def create_strategy(
+    config: TimingManagerConfig,
+    mock_credit_manager: MockCreditManager,
+) -> RequestRateStrategy:
+    """Create a RequestRateStrategy with mock dataset metadata based on config."""
+    dataset_metadata = create_mock_dataset_metadata(
+        conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
+    )
+    return RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+
+
 class TestBenchmarkDurationConfiguration:
     """Test configuration validation and behavior of benchmark duration."""
 
@@ -207,11 +218,7 @@ async def test_strategy_uses_duration_for_profiling_phase(
         """Test that RequestRateStrategy respects benchmark duration."""
         config = benchmark_duration_config(benchmark_duration=2.0)
 
-        # Create strategy and check profiling phase config
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Check that the profiling phase is configured correctly
         assert len(strategy.ordered_phase_configs) > 0
@@ -228,11 +235,7 @@ async def test_strategy_ignores_request_count_when_duration_set(
     ):
         """Test that request count is ignored when duration is specified."""
         config = mixed_config(request_count=50, benchmark_duration=1.5)
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         profiling_config = strategy.ordered_phase_configs[-1]
 
@@ -245,11 +248,7 @@ async def test_strategy_fallback_to_request_count(
     ):
         """Test that strategy falls back to request count when no duration."""
         config = mixed_config(request_count=25, benchmark_duration=None)
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         profiling_config = strategy.ordered_phase_configs[-1]
 
@@ -264,11 +263,7 @@ async def test_strategy_with_warmup_and_duration(
         config = benchmark_duration_config(
             benchmark_duration=4.0, warmup_request_count=10
         )
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Should have warmup and profiling phases
         assert len(strategy.ordered_phase_configs) == 2
@@ -294,11 +289,7 @@ async def test_strategy_with_duration_and_concurrency(
     ):
         """Test strategy with duration and concurrency settings."""
         config = benchmark_duration_config(benchmark_duration=3.0, concurrency=5)
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         assert config.benchmark_duration == 3.0
         assert config.concurrency == 5
@@ -335,11 +326,7 @@ async def test_various_duration_warmup_combinations(
         config = benchmark_duration_config(
             benchmark_duration=duration, warmup_request_count=warmup_count
         )
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Verify configuration is correct
         assert config.benchmark_duration == duration
@@ -438,10 +425,7 @@ async def test_profiling_phase_setup_with_duration(
     ):
         """Test profiling phase setup when duration is specified."""
         config = benchmark_duration_config(benchmark_duration=8.0)
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Find the profiling phase config
         profiling_config = next(
@@ -462,10 +446,7 @@ async def test_profiling_phase_setup_without_duration(
     ):
         """Test profiling phase setup when duration is not specified."""
         config = mixed_config(request_count=40, benchmark_duration=None)
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Find the profiling phase config
         profiling_config = next(
@@ -488,10 +469,7 @@ async def test_warmup_phase_unaffected_by_duration(
         config = benchmark_duration_config(
             benchmark_duration=12.0, warmup_request_count=15
         )
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Find the warmup phase config
         warmup_config = next(
@@ -646,10 +624,7 @@ async def test_force_completion_when_timeout_triggered(self, time_traveler):
         # Create a time-based phase that has already exceeded duration
         config = benchmark_duration_config(benchmark_duration=1.0)
         mock_credit_manager = MockCreditManager(time_traveler=time_traveler)
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Create a phase stats that would normally have in-flight requests
         phase_stats = CreditPhaseStats(
@@ -686,10 +661,7 @@ async def test_wait_for_phase_completion_with_timeout(self, time_traveler):
         """Test that _wait_for_phase_completion respects duration timeout."""
         config = benchmark_duration_config(benchmark_duration=2.0)
         mock_credit_manager = MockCreditManager(time_traveler=time_traveler)
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Create a time-based phase that is close to expiring
         start_time = time_traveler.time_ns()
@@ -729,10 +701,7 @@ async def test_wait_for_phase_completion_without_timeout_for_request_count(self)
         mock_credit_manager = MockCreditManager(
             time_traveler=None
         )  # No time manipulation needed
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Create a request-count-based phase
         phase_stats = CreditPhaseStats(
@@ -836,11 +805,7 @@ async def test_grace_period_with_strategy(
         config = benchmark_duration_config(
             benchmark_duration=2.0, benchmark_grace_period=15.0
         )
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         assert strategy.config.benchmark_grace_period == 15.0
         assert strategy.config.benchmark_duration == 2.0
@@ -853,11 +818,7 @@ async def test_grace_period_integration_with_duration(
             benchmark_duration=1.0,  # Short duration for testing
             benchmark_grace_period=5.0,
         )
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Should have profiling phase with duration configuration
         profiling_config = strategy.ordered_phase_configs[-1]
@@ -879,11 +840,7 @@ async def test_grace_period_with_quick_completion(
         config = benchmark_duration_config(
             benchmark_duration=1.0, benchmark_grace_period=5.0
         )
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Create a profiling phase that completes quickly
         phase_stats = CreditPhaseStats(
@@ -912,11 +869,7 @@ async def test_grace_period_timeout_with_in_flight_requests(
         config = benchmark_duration_config(
             benchmark_duration=1.0, benchmark_grace_period=2.0
         )
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Create a profiling phase with in-flight requests
         phase_stats = CreditPhaseStats(
@@ -959,11 +912,7 @@ async def test_zero_grace_period_immediate_completion(
         config = benchmark_duration_config(
             benchmark_duration=1.0, benchmark_grace_period=0.0
         )
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Create a profiling phase with in-flight requests
         phase_stats = CreditPhaseStats(
@@ -993,11 +942,7 @@ async def test_grace_period_completion_during_grace_period(
         config = benchmark_duration_config(
             benchmark_duration=1.0, benchmark_grace_period=5.0
         )
-
-        dataset_metadata = create_mock_dataset_metadata(
-            conversation_ids=[f"conv{i}" for i in range(config.request_count or 10)]
-        )
-        strategy = RequestRateStrategy(config, mock_credit_manager, dataset_metadata)
+        strategy = create_strategy(config, mock_credit_manager)
 
         # Create a profiling phase with in-flight requests
         phase_stats = CreditPhaseStats(