初始化提交

backend/tests/ai_provider/llm/test_pool.py

@@ -0,0 +1,100 @@
+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)