From d91551540e983b3818f722fb3f86be77d0b445e2 Mon Sep 17 00:00:00 2001
From: dalthecow <dalcowboiz@gmail.com>
Date: Thu, 11 Jul 2024 17:40:14 -0400
Subject: [PATCH] test sweep profile generator

---
 tests/unit/executor/test_profile_generator.py | 91 +++++++++++++++++--
 1 file changed, 82 insertions(+), 9 deletions(-)

diff --git a/tests/unit/executor/test_profile_generator.py b/tests/unit/executor/test_profile_generator.py
index 1c9f084..8ffb6cc 100644
--- a/tests/unit/executor/test_profile_generator.py
+++ b/tests/unit/executor/test_profile_generator.py
@@ -1,7 +1,9 @@
+import numpy
 import pytest
 from unittest.mock import MagicMock
 from guidellm.executor import (ProfileGenerator, FixedRateProfileGenerator, SweepProfileGenerator)
-from src.guidellm.core.result import TextGenerationBenchmarkReport
+from src.guidellm.core.result import TextGenerationBenchmark, TextGenerationBenchmarkReport
+from src.guidellm.executor import profile_generator
 from src.guidellm.scheduler.load_generator import LoadGenerationModes
 
 def test_invalid_profile_generation_mode_error():
@@ -11,13 +13,7 @@ def test_invalid_profile_generation_mode_error():
   with pytest.raises(ValueError, match=f"Invalid profile generation mode: {profile_mode}"):
     ProfileGenerator.create_generator(profile_mode, **({ "rate": rate, "rate_type": rate_type}))
 
-def test_sweep_profile_generator_creation():
-  profile_generator = ProfileGenerator.create_generator("sweep", **({}))
-  assert isinstance(profile_generator, SweepProfileGenerator)
-  assert profile_generator._sync_run == False
-  assert profile_generator._max_found == False
-  assert profile_generator._pending_rates == None
-  assert profile_generator._pending_rates == None
+# Fixed Rate Profile Generator
 
 def test_fixed_rate_profile_generator_creation():
   rate = [1]
@@ -53,4 +49,81 @@ def test_next_profile_with_sync_mode():
   current_profile = profile_generator.next_profile(mock_report)
   assert current_profile.load_gen_rate == None
   assert current_profile.load_gen_mode.name == LoadGenerationModes.SYNCHRONOUS.name
-  assert profile_generator.next_profile(mock_report) == None
\ No newline at end of file
+  assert profile_generator.next_profile(mock_report) == None
+
+# Sweep Profile Generator
+
+def test_sweep_profile_generator_creation():
+  profile_generator = ProfileGenerator.create_generator("sweep", **({}))
+  assert isinstance(profile_generator, SweepProfileGenerator)
+  assert profile_generator._sync_run == False
+  assert profile_generator._max_found == False
+  assert profile_generator._pending_rates == None
+  assert profile_generator._pending_rates == None
+
+def test_first_profile_is_synchronous():
+  profile_generator = ProfileGenerator.create_generator("sweep")
+  mock_report = MagicMock(spec=TextGenerationBenchmarkReport)
+  profile = profile_generator.next_profile(mock_report)
+  assert profile.load_gen_rate == None
+  assert profile.load_gen_mode.name == LoadGenerationModes.SYNCHRONOUS.name
+
+def test_rate_doubles():
+  profile_generator = ProfileGenerator.create_generator("sweep")
+  mock_report = MagicMock(spec=TextGenerationBenchmarkReport)
+  mock_benchmark = MagicMock(spec=TextGenerationBenchmark)
+  mock_benchmark.overloaded = False
+  mock_benchmark.args_rate = 2.0
+  mock_benchmark.request_rate = 2.0
+  benchmarks = [
+        mock_benchmark
+    ]
+  mock_report.benchmarks = benchmarks
+  profile = profile_generator.next_profile(mock_report)
+
+  profile = profile_generator.next_profile(mock_report)
+  assert profile.load_gen_rate == 4.0
+
+def test_max_found():
+  profile_generator = ProfileGenerator.create_generator("sweep")
+  mock_report = MagicMock(spec=TextGenerationBenchmarkReport)
+  mock_benchmark = MagicMock(spec=TextGenerationBenchmark)
+  mock_benchmark.overloaded = False
+  mock_benchmark.args_rate = 2.0
+  mock_benchmark.request_rate = 2.0
+  mock_overloaded_benchmark = MagicMock(spec=TextGenerationBenchmark)
+  mock_overloaded_benchmark.overloaded = True
+  mock_overloaded_benchmark.args_rate = 4.0
+  mock_overloaded_benchmark.request_rate = 4.0
+  benchmarks = [
+        mock_benchmark,
+        mock_overloaded_benchmark
+    ]
+  mock_report.benchmarks = benchmarks
+
+  profile_generator.next_profile(mock_report)
+  profile = profile_generator.next_profile(mock_report)
+
+  # if benchmark wasn't overloaded, rate would have doubled to 8
+  assert profile.load_gen_rate == 2.0
+
+def test_pending_rates():
+  profile_generator = ProfileGenerator.create_generator("sweep")
+  mock_report = MagicMock(spec=TextGenerationBenchmarkReport)
+  mock_benchmark = MagicMock(spec=TextGenerationBenchmark)
+  mock_benchmark.overloaded = False
+  mock_benchmark.args_rate = 2.0
+  mock_benchmark.request_rate = 2.0
+  mock_overloaded_benchmark = MagicMock(spec=TextGenerationBenchmark)
+  mock_overloaded_benchmark.overloaded = True
+  mock_overloaded_benchmark.args_rate = 8.0
+  mock_overloaded_benchmark.request_rate = 8.0
+  benchmarks = [
+        mock_benchmark,
+        mock_overloaded_benchmark
+    ]
+  mock_report.benchmarks = benchmarks
+  profile = profile_generator.next_profile(mock_report)
+  for expected_rate in numpy.linspace(2.0, 8.0, 10):
+        profile = profile_generator.next_profile(mock_report)
+        assert profile.load_gen_rate == expected_rate
\ No newline at end of file