From 2aa0f80484a1497f14ded746d31a38818c2358e1 Mon Sep 17 00:00:00 2001
From: Daisong Gan <81140776+Hubuguilai@users.noreply.github.com>
Date: Sat, 26 Sep 2026 20:37:20 +0800
Subject: [PATCH] test: cover the zero-budget tag check and tidy the regression
 test

Add the tag-budget case (the third changed site was previously untested),
drop the comments that only restated the assertions, and give each case its
own limiter and result names instead of rebinding the same locals.
---
 .../router_strategy/test_budget_limiter.py    | 37 +++++++++++++------
 1 file changed, 25 insertions(+), 12 deletions(-)

diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py
index 1eaba929218..44b6c585806 100644
--- a/tests/unit/router_strategy/test_budget_limiter.py
+++ b/tests/unit/router_strategy/test_budget_limiter.py
@@ -154,13 +154,12 @@ async def test_zero_max_budget_blocks_spend(disable_budget_sync):
         }
     ]
 
-    # Provider budget of 0 blocks a deployment that already spent.
-    limiter = RouterBudgetLimiting(
+    provider_limiter = RouterBudgetLimiting(
         dual_cache=DualCache(),
         provider_budget_config={"openai": {"budget_limit": 0.0, "time_period": "1d"}},
     )
     provider_configs: Final = {"openai": BudgetConfig(max_budget=0.0, budget_duration="1d")}
-    kept, _ = limiter._filter_out_deployments_above_budget(
+    provider_kept, _ = provider_limiter._filter_out_deployments_above_budget(
         potential_deployments=[],
         healthy_deployments=healthy_deployments,
         provider_configs=provider_configs,
@@ -169,10 +168,9 @@ async def test_zero_max_budget_blocks_spend(disable_budget_sync):
         spend_map={"provider_spend:openai:1d": 5.0},
         request_tags=[],
     )
-    assert kept == []
+    assert provider_kept == []
 
-    # Deployment budget of 0 blocks even at zero spend.
-    limiter = RouterBudgetLimiting(
+    deployment_limiter = RouterBudgetLimiting(
         dual_cache=DualCache(),
         provider_budget_config=None,
         model_list=[
@@ -184,7 +182,7 @@ async def test_zero_max_budget_blocks_spend(disable_budget_sync):
         ],
     )
     deployment_configs: Final = {"deployment-1": BudgetConfig(max_budget=0.0, budget_duration="1d")}
-    kept, _ = limiter._filter_out_deployments_above_budget(
+    deployment_kept, _ = deployment_limiter._filter_out_deployments_above_budget(
         potential_deployments=[],
         healthy_deployments=healthy_deployments,
         provider_configs={},
@@ -193,14 +191,29 @@ async def test_zero_max_budget_blocks_spend(disable_budget_sync):
         spend_map={"deployment_spend:deployment-1:1d": 0.0},
         request_tags=[],
     )
-    assert kept == []
+    assert deployment_kept == []
 
-    # A None budget still means "no limit" and must stay routable.
-    limiter = RouterBudgetLimiting(
+    tag_limiter = RouterBudgetLimiting(
         dual_cache=DualCache(),
         provider_budget_config=None,
     )
-    kept, _ = limiter._filter_out_deployments_above_budget(
+    tag_limiter.tag_budget_config = {"blocked-tag": BudgetConfig(time_period="1d", budget_limit=0.0)}
+    tag_kept, _ = tag_limiter._filter_out_deployments_above_budget(
+        potential_deployments=[],
+        healthy_deployments=healthy_deployments,
+        provider_configs={},
+        deployment_configs={},
+        deployment_providers=["openai"],
+        spend_map={"tag_spend:blocked-tag:1d": 0.0},
+        request_tags=["blocked-tag"],
+    )
+    assert tag_kept == []
+
+    unset_limiter = RouterBudgetLimiting(
+        dual_cache=DualCache(),
+        provider_budget_config=None,
+    )
+    unset_kept, _ = unset_limiter._filter_out_deployments_above_budget(
         potential_deployments=[],
         healthy_deployments=healthy_deployments,
         provider_configs={},
@@ -209,4 +222,4 @@ async def test_zero_max_budget_blocks_spend(disable_budget_sync):
         spend_map={},
         request_tags=[],
     )
-    assert len(kept) == 1
+    assert len(unset_kept) == 1