test(budget): add budget management coverage#284
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
💤 Files with no reviewable changes (1)
📝 WalkthroughWalkthroughAdds extensive tests and test helpers for budget enforcement across unit, SQLite, integration (Postgres/Mongo), contract, and E2E suites; wires budget checker and pricing resolver into test servers and updates test storage assertions and usage cost fields. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Client
participant Server as E2E Server
participant Budget as BudgetChecker
participant Store as UsageStore
participant Provider as UpstreamProvider
participant Audit as AuditLog
Client->>Server: POST /chat (user-path, req-id, payload)
Server->>Budget: Check(user-path, req-id, costEstimate)
Budget->>Store: SumUsageCost(user-path, periodStart, now)
Store-->>Budget: currentSpent
alt currentSpent + costEstimate < budget.Amount
Budget-->>Server: Allowed
Server->>Provider: Forward request
Provider-->>Server: Response + usage
Server->>Store: Persist usage entry
Server->>Audit: Write audit (if enabled)
Server-->>Client: 200 OK + response
else currentSpent + costEstimate >= budget.Amount
Budget-->>Server: ExceededError (spent)
Server->>Audit: Write audit with rate_limit (if enabled)
Server-->>Client: 429 Too Many Requests (Retry-After, budget_exceeded)
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR adds comprehensive budget management test coverage across unit, component E2E, integration, contract, and release curl scenarios. It also rebases a test-only branch onto All changes are test-only and the test logic is generally well-structured, with deterministic boundary cases, proper async polling patterns, and correct per-fixture DB isolation. Confidence Score: 5/5Safe to merge — all changes are test-only and the single finding is a minor semantic mismatch that does not affect test correctness in practice. The only finding is a P2 style issue ( tests/e2e/budget_test.go – Important Files Changed
Sequence DiagramsequenceDiagram
participant T as Test
participant S as Gateway Server
participant B as BudgetService
participant U as UsageLogger
participant DB as SQLite/PG/Mongo
T->>S: POST /v1/chat/completions (req1)
S->>B: Check(userPath, now)
B->>DB: SumUsageCost(userPath, periodStart, now)
DB-->>B: (0, false, nil) — no usage yet
B-->>S: nil — allow
S->>T: 200 OK
S->>U: LogUsage(req1, cost)
U-->>DB: flush after 10ms
T->>T: waitForBudgetSpent polls Statuses() until Spent >= amount
T->>S: POST /v1/chat/completions (req2)
S->>B: Check(userPath, now)
B->>DB: SumUsageCost(userPath, periodStart, now)
DB-->>B: (0.018, true, nil) — usage flushed
B-->>S: ExceededError (spent >= amount)
S->>T: 429 Too Many Requests + Retry-After header
Reviews (1): Last reviewed commit: "test(budget): add budget management cove..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if err != nil || len(statuses) != 1 { | ||
| return false | ||
| } | ||
| return statuses[0].Budget.UserPath == userPath && statuses[0].HasUsage && statuses[0].Spent > amount |
There was a problem hiding this comment.
waitForBudgetSpent polling condition uses strict > instead of >=
The enforcement threshold is spent >= amount (as verified by TestServiceCheckBudgetAmountBoundary's "equal amount blocks" case), but the polling here uses >. If the first request's cost lands exactly on the budget amount, this helper will not return, timing out after 2 seconds even though the budget is already enforced. In practice the mock always produces cost > amount, but the condition is semantically misleading.
| return statuses[0].Budget.UserPath == userPath && statuses[0].HasUsage && statuses[0].Spent > amount | |
| return statuses[0].Budget.UserPath == userPath && statuses[0].HasUsage && statuses[0].Spent >= amount |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/e2e/budget_test.go`:
- Around line 168-178: The close method currently captures only the first error
from usageLogger.Close() and discards db.Close()'s error; change f.closeErr
assignment inside f.closeOnce.Do to join both errors using errors.Join so both
usageLogger.Close() and db.Close() failures are preserved; update the block
handling f.closeErr to set f.closeErr = errors.Join(f.closeErr, err) (importing
the "errors" package) while keeping the existing require.NoError(t, f.closeErr)
behavior so CI logs show both errors.
In `@tests/e2e/release-e2e-scenarios.md`:
- Around line 1488-1491: The test currently hardcodes a trailing PUT to
"$BASE_URL/admin/api/v1/budgets/settings" with fixed values (daily_reset_hour:0,
weekly_reset_weekday:1, monthly_reset_day:1, etc.); change it to capture and
persist the pre-test settings before running S86 and then restore those exact
values after S86 completes by PUTting the saved payload back to the same
endpoint. Specifically, add a read GET to
"$BASE_URL/admin/api/v1/budgets/settings" to save the returned JSON, and replace
the hardcoded body in the final PUT with the saved JSON values for keys
daily_reset_hour, daily_reset_minute, weekly_reset_weekday, weekly_reset_hour,
weekly_reset_minute, monthly_reset_day, monthly_reset_hour, and
monthly_reset_minute so the environment is restored to its original state.
In `@tests/integration/budget_test.go`:
- Around line 134-151: The t.Fatalf inside the require.Eventually callback runs
on a worker goroutine and will not report properly; fix by resolving the DB-type
dispatch before calling require.Eventually: in waitForAuditEntry determine a
fetch function or closure based on fixture.DBType (call
dbassert.QueryAuditLogsByRequestID or dbassert.QueryAuditLogsByRequestIDMongo)
and call t.Fatalf immediately for an unsupported DB type, then use that fetcher
inside the require.Eventually callback; apply the same pattern to
waitForBudgetUsage / usageEntriesByRequestID so no t.Fatalf executes from within
the Eventually worker goroutine.
In `@tests/integration/dbassert/usage.go`:
- Around line 238-246: The cost decoding only checks for float64 when reading
doc["input_cost"], doc["output_cost"], and doc["total_cost"], so whole-number
BSON integers can leave those pointers nil; update the decode logic (the blocks
that set entry.InputCost, entry.OutputCost, entry.TotalCost) to also accept
int32 and int64 by converting them to float64 (or factor the logic into a helper
like bsonToFloat64(v any) (float64, bool)) and use its result to set the
pointers when present.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: bae5b06a-d963-40f4-b655-ed4d2706aa26
📒 Files selected for processing (11)
internal/budget/service_test.gointernal/budget/store_sqlite_test.gotests/contract/usage_contract_test.gotests/e2e/budget_test.gotests/e2e/manage-release-e2e-stack.shtests/e2e/release-e2e-scenarios.mdtests/e2e/setup_test.gotests/integration/budget_test.gotests/integration/dbassert/budget.gotests/integration/dbassert/usage.gotests/integration/setup_test.go
| func (f *budgetE2EFixture) close(t *testing.T) { | ||
| t.Helper() | ||
|
|
||
| f.closeOnce.Do(func() { | ||
| f.closeErr = f.usageLogger.Close() | ||
| if err := f.db.Close(); err != nil && f.closeErr == nil { | ||
| f.closeErr = err | ||
| } | ||
| }) | ||
| require.NoError(t, f.closeErr) | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Close-error precedence drops the secondary error silently.
If usageLogger.Close() fails, db.Close()'s error is intentionally swallowed (first-error-wins). This is a common pattern but means a leaked SQLite handle wouldn't be visible alongside a logger flush failure. Acceptable for tests; flagging only because both are real resources and either failure is interesting in CI logs.
♻️ Surface both close errors via errors.Join
f.closeOnce.Do(func() {
- f.closeErr = f.usageLogger.Close()
- if err := f.db.Close(); err != nil && f.closeErr == nil {
- f.closeErr = err
- }
+ f.closeErr = errors.Join(f.usageLogger.Close(), f.db.Close())
})(Requires importing "errors".)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/e2e/budget_test.go` around lines 168 - 178, The close method currently
captures only the first error from usageLogger.Close() and discards db.Close()'s
error; change f.closeErr assignment inside f.closeOnce.Do to join both errors
using errors.Join so both usageLogger.Close() and db.Close() failures are
preserved; update the block handling f.closeErr to set f.closeErr =
errors.Join(f.closeErr, err) (importing the "errors" package) while keeping the
existing require.NoError(t, f.closeErr) behavior so CI logs show both errors.
| curl -fsS -X PUT "$BASE_URL/admin/api/v1/budgets/settings" \ | ||
| -H 'Content-Type: application/json' \ | ||
| -d '{"daily_reset_hour":0,"daily_reset_minute":0,"weekly_reset_weekday":1,"weekly_reset_hour":0,"weekly_reset_minute":0,"monthly_reset_day":1,"monthly_reset_hour":0,"monthly_reset_minute":0}' \ | ||
| >/dev/null |
There was a problem hiding this comment.
Settings "reset" hardcodes assumed defaults rather than restoring pre-test state.
The trailing PUT writes daily_reset_hour:0, weekly_reset_weekday:1, monthly_reset_day:1, etc., which is a fixed known state, not the values that were in place before S86 ran. If the system's actual defaults ever change, this will silently leave the gateway in a non-default configuration after S86. Since later scenarios (S87–S89) only touch budgets and not settings, this is not a functional issue today, but worth flagging if you want this matrix to be fully idempotent.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/e2e/release-e2e-scenarios.md` around lines 1488 - 1491, The test
currently hardcodes a trailing PUT to "$BASE_URL/admin/api/v1/budgets/settings"
with fixed values (daily_reset_hour:0, weekly_reset_weekday:1,
monthly_reset_day:1, etc.); change it to capture and persist the pre-test
settings before running S86 and then restore those exact values after S86
completes by PUTting the saved payload back to the same endpoint. Specifically,
add a read GET to "$BASE_URL/admin/api/v1/budgets/settings" to save the returned
JSON, and replace the hardcoded body in the final PUT with the saved JSON values
for keys daily_reset_hour, daily_reset_minute, weekly_reset_weekday,
weekly_reset_hour, weekly_reset_minute, monthly_reset_day, monthly_reset_hour,
and monthly_reset_minute so the environment is restored to its original state.
| func waitForAuditEntry(t *testing.T, fixture *TestServerFixture, requestID string) dbassert.AuditLogEntry { | ||
| t.Helper() | ||
|
|
||
| var entries []dbassert.AuditLogEntry | ||
| require.Eventually(t, func() bool { | ||
| switch fixture.DBType { | ||
| case "postgresql": | ||
| entries = dbassert.QueryAuditLogsByRequestID(t, fixture.PgPool, requestID) | ||
| case "mongodb": | ||
| entries = dbassert.QueryAuditLogsByRequestIDMongo(t, fixture.MongoDb, requestID) | ||
| default: | ||
| t.Fatalf("unsupported DB type %q", fixture.DBType) | ||
| } | ||
| return len(entries) == 1 | ||
| }, 5*time.Second, 100*time.Millisecond) | ||
|
|
||
| return entries[0] | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
testify require.Eventually goroutine t.Fatal
💡 Result:
Do not use require.Eventually(t, ...) with require.* assertions inside the condition function, as the condition runs in a separate goroutine spawned by assert.Eventually (which require.Eventually forwards to). require.* calls t.FailNow, which is unsafe from non-main test goroutines per Go's testing package docs and can cause panics, hangs, or flaky tests (e.g., goroutine leaks if condition overruns timeout). Use assert.Eventually or assert.EventuallyWithT instead inside the condition: - For simple bool conditions: assert.Eventually(t, func bool { ... }, waitFor, tick) - For multiple assertions: assert.EventuallyWithT(t, func(c assert.CollectT) { assert.NoError(c, err); ... }, waitFor, tick) - Use assert.(c, ...) or require.(c, ...) with the provided assert.CollectT; these collect failures for retry until final tick. - Do NOT pass outer t to require.(t, ...) in condition. require.Eventually(t, ...) is safe for the top-level call (runs in main goroutine) but problematic inside condition due to goroutine. Known issues: goroutine leaks (#1611), hangs on FailNow/Goexit (#1810), panics post-completion. Use assert versions to avoid. Examples: require.Eventually(t, func bool { return len(queue) == 0 // simple bool, no assertions }, time.Second, 10time.Millisecond) // Better for assertions: assert.EventuallyWithT(t, func(c assert.CollectT) { items, err := fetch require.NoError(c, err) // safe with c assert.Len(c, items, 0) }, time.Second, 10time.Millisecond) pkg.go.dev docs warn: require.* must be from main test goroutine.
Citations:
- 1: Detect and fail on
require.*methods called from a wrong goroutine stretchr/testify#1499 - 2: Improve the behavior of
requireassertions when usingassert.EventuallyWithTstretchr/testify#1396 - 3:
Eventuallyhangs on failed/abortedconditionstretchr/testify#1810 - 4: How to handle failed expectations inside of a goroutine? stretchr/testify#772
- 5: https://pkg.go.dev/github.com/stretchr/testify/require?utm%5C_source=godoc
- 6: https://pkg.go.dev/mod/github.com/stretchr/testify
- 7: https://github.com/stretchr/testify/tree/master/require
- 8: https://github.com/stretchr/testify/blob/master/require/require.go
🏁 Script executed:
# Find and examine the test file
fd budget_test.goRepository: ENTERPILOT/GoModel
Length of output: 120
🏁 Script executed:
# Read the relevant sections of the test file
head -160 tests/integration/budget_test.go | tail -60Repository: ENTERPILOT/GoModel
Length of output: 1988
🏁 Script executed:
# Also check the beginning to see the test table
head -50 tests/integration/budget_test.goRepository: ENTERPILOT/GoModel
Length of output: 1411
t.Fatalf inside require.Eventually callback runs on a worker goroutine.
require.Eventually invokes the condition function in a separate goroutine. The Go testing package requires Fatal/Fatalf to be called only from the test goroutine; from a worker goroutine runtime.Goexit only kills the worker, so Eventually keeps ticking until the 5s timeout and you lose the error message. The same concern applies transitively in waitForBudgetUsage (line 113) via usageEntriesByRequestID's t.Fatalf (line 129).
In practice the table-driven test only passes "postgresql"/"mongodb" so the default branches never trigger, but this pattern will silently degrade error reporting if DB type validation or fallback logic ever changes.
♻️ Resolve the DB selection once, outside Eventually
func waitForAuditEntry(t *testing.T, fixture *TestServerFixture, requestID string) dbassert.AuditLogEntry {
t.Helper()
+ var query func() []dbassert.AuditLogEntry
+ switch fixture.DBType {
+ case "postgresql":
+ query = func() []dbassert.AuditLogEntry {
+ return dbassert.QueryAuditLogsByRequestID(t, fixture.PgPool, requestID)
+ }
+ case "mongodb":
+ query = func() []dbassert.AuditLogEntry {
+ return dbassert.QueryAuditLogsByRequestIDMongo(t, fixture.MongoDb, requestID)
+ }
+ default:
+ t.Fatalf("unsupported DB type %q", fixture.DBType)
+ }
+
var entries []dbassert.AuditLogEntry
require.Eventually(t, func() bool {
- switch fixture.DBType {
- case "postgresql":
- entries = dbassert.QueryAuditLogsByRequestID(t, fixture.PgPool, requestID)
- case "mongodb":
- entries = dbassert.QueryAuditLogsByRequestIDMongo(t, fixture.MongoDb, requestID)
- default:
- t.Fatalf("unsupported DB type %q", fixture.DBType)
- }
+ entries = query()
return len(entries) == 1
}, 5*time.Second, 100*time.Millisecond)
return entries[0]
}A similar treatment for waitForBudgetUsage resolves the indirect t.Fatalf reached via usageEntriesByRequestID: move the DB-type dispatch into a closure or helper, and call it from the Eventually callback to ensure the error check runs in the test goroutine.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func waitForAuditEntry(t *testing.T, fixture *TestServerFixture, requestID string) dbassert.AuditLogEntry { | |
| t.Helper() | |
| var entries []dbassert.AuditLogEntry | |
| require.Eventually(t, func() bool { | |
| switch fixture.DBType { | |
| case "postgresql": | |
| entries = dbassert.QueryAuditLogsByRequestID(t, fixture.PgPool, requestID) | |
| case "mongodb": | |
| entries = dbassert.QueryAuditLogsByRequestIDMongo(t, fixture.MongoDb, requestID) | |
| default: | |
| t.Fatalf("unsupported DB type %q", fixture.DBType) | |
| } | |
| return len(entries) == 1 | |
| }, 5*time.Second, 100*time.Millisecond) | |
| return entries[0] | |
| } | |
| func waitForAuditEntry(t *testing.T, fixture *TestServerFixture, requestID string) dbassert.AuditLogEntry { | |
| t.Helper() | |
| var query func() []dbassert.AuditLogEntry | |
| switch fixture.DBType { | |
| case "postgresql": | |
| query = func() []dbassert.AuditLogEntry { | |
| return dbassert.QueryAuditLogsByRequestID(t, fixture.PgPool, requestID) | |
| } | |
| case "mongodb": | |
| query = func() []dbassert.AuditLogEntry { | |
| return dbassert.QueryAuditLogsByRequestIDMongo(t, fixture.MongoDb, requestID) | |
| } | |
| default: | |
| t.Fatalf("unsupported DB type %q", fixture.DBType) | |
| } | |
| var entries []dbassert.AuditLogEntry | |
| require.Eventually(t, func() bool { | |
| entries = query() | |
| return len(entries) == 1 | |
| }, 5*time.Second, 100*time.Millisecond) | |
| return entries[0] | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/integration/budget_test.go` around lines 134 - 151, The t.Fatalf inside
the require.Eventually callback runs on a worker goroutine and will not report
properly; fix by resolving the DB-type dispatch before calling
require.Eventually: in waitForAuditEntry determine a fetch function or closure
based on fixture.DBType (call dbassert.QueryAuditLogsByRequestID or
dbassert.QueryAuditLogsByRequestIDMongo) and call t.Fatalf immediately for an
unsupported DB type, then use that fetcher inside the require.Eventually
callback; apply the same pattern to waitForBudgetUsage / usageEntriesByRequestID
so no t.Fatalf executes from within the Eventually worker goroutine.
| if v, ok := doc["input_cost"].(float64); ok { | ||
| entry.InputCost = &v | ||
| } | ||
| if v, ok := doc["output_cost"].(float64); ok { | ||
| entry.OutputCost = &v | ||
| } | ||
| if v, ok := doc["total_cost"].(float64); ok { | ||
| entry.TotalCost = &v | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider int32/int64 fallback for cost decoding for parity with token decoding.
Token fields above (lines 223-237) fall back through int32/int64 because BSON encodes whole numbers as integer types. Cost fields here only check float64, so a document where a cost happens to be persisted as a whole number (e.g., a fixture written with bson.M{"input_cost": 0}) would silently leave the pointer nil and surface as a hard-to-debug assertion failure in tests/integration/budget_test.go (which waits for non-nil TotalCost). In normal flows the producer writes *float64, so this is unlikely to occur in production data, but adding the fallback would make this assertion helper more robust.
♻️ Optional refactor
- if v, ok := doc["input_cost"].(float64); ok {
- entry.InputCost = &v
- }
- if v, ok := doc["output_cost"].(float64); ok {
- entry.OutputCost = &v
- }
- if v, ok := doc["total_cost"].(float64); ok {
- entry.TotalCost = &v
- }
+ if v, ok := bsonToFloat64(doc["input_cost"]); ok {
+ entry.InputCost = &v
+ }
+ if v, ok := bsonToFloat64(doc["output_cost"]); ok {
+ entry.OutputCost = &v
+ }
+ if v, ok := bsonToFloat64(doc["total_cost"]); ok {
+ entry.TotalCost = &v
+ }func bsonToFloat64(v any) (float64, bool) {
switch n := v.(type) {
case float64:
return n, true
case int32:
return float64(n), true
case int64:
return float64(n), true
}
return 0, false
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/integration/dbassert/usage.go` around lines 238 - 246, The cost
decoding only checks for float64 when reading doc["input_cost"],
doc["output_cost"], and doc["total_cost"], so whole-number BSON integers can
leave those pointers nil; update the decode logic (the blocks that set
entry.InputCost, entry.OutputCost, entry.TotalCost) to also accept int32 and
int64 by converting them to float64 (or factor the logic into a helper like
bsonToFloat64(v any) (float64, bool)) and use its result to set the pointers
when present.
2 participants
Summary
mainafter the budget management base PR was mergedCHANGELOG.md; it was unreferenced and release notes are handled through the PR/release workflowTests
go test -count=1 ./internal/budgetgo test -count=1 -tags=e2e ./tests/e2ego test -count=1 -tags=contract ./tests/contractgo test -tags=integration ./tests/integrationbash -n tests/e2e/manage-release-e2e-stack.sh tests/e2e/run-release-e2e.shSummary by CodeRabbit
Tests
Documentation