101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
from app.ai_provider.llm.pool import (
|
|
CostAwareStrategy,
|
|
EndpointPool,
|
|
LeastConnectionsStrategy,
|
|
LowestLatencyStrategy,
|
|
PoolFactory,
|
|
PriorityFailoverStrategy,
|
|
WeightedRoundRobinStrategy,
|
|
)
|
|
|
|
|
|
class TestWeightedRoundRobin:
|
|
def test_distribution(self):
|
|
endpoints = [{"name": "a", "weight": 100}, {"name": "b", "weight": 50}]
|
|
strategy = WeightedRoundRobinStrategy(endpoints)
|
|
picks = [strategy.select([])["name"] for _ in range(15)]
|
|
assert picks.count("a") > picks.count("b")
|
|
|
|
def test_skip_excluded(self):
|
|
endpoints = [{"name": "a", "weight": 100}, {"name": "b", "weight": 100}]
|
|
strategy = WeightedRoundRobinStrategy(endpoints)
|
|
pick = strategy.select(["a"])
|
|
assert pick["name"] == "b"
|
|
|
|
|
|
class TestPriorityFailover:
|
|
def test_primary_selected(self):
|
|
endpoints = [{"name": "primary", "priority": 10}, {"name": "backup", "priority": 20}]
|
|
strategy = PriorityFailoverStrategy(endpoints)
|
|
pick = strategy.select([])
|
|
assert pick["name"] == "primary"
|
|
|
|
def test_failover_on_excluded(self):
|
|
endpoints = [{"name": "primary", "priority": 10}, {"name": "backup", "priority": 20}]
|
|
strategy = PriorityFailoverStrategy(endpoints)
|
|
pick = strategy.select(["primary"])
|
|
assert pick["name"] == "backup"
|
|
|
|
def test_exhausted(self):
|
|
endpoints = [{"name": "a", "priority": 10}]
|
|
strategy = PriorityFailoverStrategy(endpoints)
|
|
pick = strategy.select(["a"])
|
|
assert pick is None
|
|
|
|
|
|
class TestLowestLatency:
|
|
def test_pick_lowest_latency(self):
|
|
endpoints = [{"name": "fast", "latency_ms": 100}, {"name": "slow", "latency_ms": 500}]
|
|
strategy = LowestLatencyStrategy(endpoints)
|
|
pick = strategy.select([])
|
|
assert pick["name"] == "fast"
|
|
|
|
|
|
class TestCostAware:
|
|
def test_pick_cheapest(self):
|
|
endpoints = [{"name": "cheap"}, {"name": "expensive"}]
|
|
|
|
def _cost(name: str) -> float:
|
|
return 0.1 if name == "cheap" else 1.0
|
|
|
|
strategy = CostAwareStrategy(endpoints, get_cost=_cost)
|
|
pick = strategy.select([])
|
|
assert pick["name"] == "cheap"
|
|
|
|
|
|
class TestLeastConnections:
|
|
def test_pick_least_connections(self):
|
|
endpoints = [{"name": "busy", "connections": 5}, {"name": "idle", "connections": 0}]
|
|
strategy = LeastConnectionsStrategy(endpoints)
|
|
pick = strategy.select([])
|
|
assert pick["name"] == "idle"
|
|
|
|
|
|
class TestEndpointPool:
|
|
def test_select_with_weighted_round_robin(self):
|
|
pool = EndpointPool(
|
|
logical_model_name="test",
|
|
endpoints_configs=[
|
|
{"name": "a", "weight": 100, "status": "active"},
|
|
{"name": "b", "weight": 50, "status": "active"},
|
|
],
|
|
strategy="weighted_round_robin",
|
|
)
|
|
selected = pool.select([])
|
|
assert selected["name"] in ("a", "b")
|
|
|
|
def test_select_all_excluded_returns_none(self):
|
|
pool = EndpointPool(
|
|
logical_model_name="test",
|
|
endpoints_configs=[{"name": "a", "weight": 100, "status": "active"}],
|
|
)
|
|
assert pool.select(["a"]) is None
|
|
|
|
def test_strategy_factory(self):
|
|
pool = PoolFactory.create_pool(
|
|
logical_model_name="test",
|
|
endpoints=[{"name": "a", "weight": 100, "status": "active"}],
|
|
strategy_name="priority_failover",
|
|
)
|
|
assert isinstance(pool.strategy, PriorityFailoverStrategy)
|