60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
import unittest
|
|
|
|
from llamacpp_ha.policies import BackendCandidate, RoundRobinPolicy
|
|
|
|
|
|
def _cands(n: int) -> list[BackendCandidate]:
|
|
return [BackendCandidate(url=f"http://b{i}", index=i) for i in range(n)]
|
|
|
|
|
|
class TestRoundRobinPolicy(unittest.TestCase):
|
|
def test_single_backend_always_chosen(self):
|
|
policy = RoundRobinPolicy()
|
|
cands = _cands(1)
|
|
for _ in range(5):
|
|
result = policy.select("m", cands)
|
|
self.assertEqual(result.url, "http://b0")
|
|
|
|
def test_distributes_evenly_across_two(self):
|
|
policy = RoundRobinPolicy()
|
|
cands = _cands(2)
|
|
chosen = [policy.select("m", cands).url for _ in range(6)]
|
|
self.assertEqual(chosen.count("http://b0"), 3)
|
|
self.assertEqual(chosen.count("http://b1"), 3)
|
|
|
|
def test_distributes_evenly_across_three(self):
|
|
policy = RoundRobinPolicy()
|
|
cands = _cands(3)
|
|
chosen = [policy.select("m", cands).url for _ in range(9)]
|
|
for i in range(3):
|
|
self.assertEqual(chosen.count(f"http://b{i}"), 3)
|
|
|
|
def test_per_model_counters_independent(self):
|
|
policy = RoundRobinPolicy()
|
|
cands = _cands(2)
|
|
r1 = policy.select("model-a", cands)
|
|
r2 = policy.select("model-b", cands)
|
|
# Both start at 0 -> both get b0
|
|
self.assertEqual(r1.url, "http://b0")
|
|
self.assertEqual(r2.url, "http://b0")
|
|
# Next calls for each model independently advance
|
|
r3 = policy.select("model-a", cands)
|
|
r4 = policy.select("model-b", cands)
|
|
self.assertEqual(r3.url, "http://b1")
|
|
self.assertEqual(r4.url, "http://b1")
|
|
|
|
def test_empty_raises(self):
|
|
policy = RoundRobinPolicy()
|
|
with self.assertRaises(ValueError):
|
|
policy.select("m", [])
|
|
|
|
def test_wraps_around(self):
|
|
policy = RoundRobinPolicy()
|
|
cands = _cands(2)
|
|
urls = [policy.select("m", cands).url for _ in range(5)]
|
|
self.assertEqual(urls, ["http://b0", "http://b1", "http://b0", "http://b1", "http://b0"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|