|
3 | 3 |
|
4 | 4 | import pytest |
5 | 5 |
|
6 | | -from floe_guard import BudgetAdvisory, BudgetGuard |
| 6 | +from floe_guard import BudgetAdvisory, BudgetGuard, ManualPrice |
7 | 7 |
|
8 | 8 |
|
9 | 9 | def test_fresh_guard_is_far_from_limit() -> None: |
@@ -66,20 +66,21 @@ def test_expected_cost_unknown_before_first_call() -> None: |
66 | 66 |
|
67 | 67 | def test_est_calls_remaining_after_a_call() -> None: |
68 | 68 | g = BudgetGuard(limit_usd=1.00) |
69 | | - g._last_llm_cost = 0.10 # warm estimate; remaining is 1.00 |
| 69 | + # Drive state through the public API (manual price for determinism), not by |
| 70 | + # poking private fields: this also proves record() feeds advisory(). |
| 71 | + g.record("m", 1_000, 0, price=ManualPrice(1e-4, 0.0)) # $0.10 spent, $0.90 left |
70 | 72 | a = g.advisory() |
71 | 73 | assert a.expected_cost == pytest.approx(0.10) |
72 | | - assert a.est_calls_remaining == 10 # floor(1.00 / 0.10) |
| 74 | + assert a.est_calls_remaining == 9 # floor(0.90 / 0.10) |
73 | 75 |
|
74 | 76 |
|
75 | 77 | def test_est_calls_remaining_uses_costlier_of_llm_and_tool() -> None: |
76 | 78 | g = BudgetGuard(limit_usd=1.00) |
77 | | - g._last_llm_cost = 0.05 |
78 | | - g._last_tool_cost = 0.20 # costlier side wins (conservative) |
79 | | - g.spent_usd = 0.40 # remaining 0.60 |
80 | | - a = g.advisory() |
81 | | - assert a.expected_cost == pytest.approx(0.20) |
82 | | - assert a.est_calls_remaining == 3 # floor(0.60 / 0.20) |
| 79 | + g.record("m", 1_000, 0, price=ManualPrice(2e-4, 0.0)) # $0.20 LLM (costlier) |
| 80 | + g.record_tool("db.query", 0.05) # $0.05 tool (cheaper) |
| 81 | + a = g.advisory() # spent 0.25, remaining 0.75 |
| 82 | + assert a.expected_cost == pytest.approx(0.20) # max(llm, tool) wins |
| 83 | + assert a.est_calls_remaining == 3 # floor(0.75 / 0.20) |
83 | 84 |
|
84 | 85 |
|
85 | 86 | def test_invalid_near_limit_bps_rejected() -> None: |
|
0 commit comments