feat(#541): add dedicated GET/PATCH /workspaces/:id/budget endpoints
- New BudgetHandler with GetBudget and PatchBudget methods
- GET returns budget_limit (null or int64 USD cents), monthly_spend,
and computed budget_remaining (null when no limit, can be negative
when over-budget so callers can see the magnitude of the overage)
- PATCH accepts {budget_limit: int64|null}; null clears the ceiling;
validates non-negative values; re-reads DB to echo final state
- Both handlers are wired in router.go under the WorkspaceAuth group
- 14 unit tests covering happy paths, 404, 400 validation, DB errors,
over-budget state, zero limit, and clear-limit round-trip
- All 20 packages pass go test ./... and go build ./... is clean
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
f8106b35be
commit
22af070ef3
171
platform/internal/handlers/budget.go
Normal file
171
platform/internal/handlers/budget.go
Normal file
@ -0,0 +1,171 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"log"
|
||||
"net/http"
|
||||
|
||||
"github.com/Molecule-AI/molecule-monorepo/platform/internal/db"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// BudgetHandler exposes per-workspace budget read/write endpoints.
|
||||
// Routes (all behind WorkspaceAuth middleware):
|
||||
//
|
||||
// GET /workspaces/:id/budget — current budget_limit, monthly_spend, budget_remaining
|
||||
// PATCH /workspaces/:id/budget — set or clear budget_limit
|
||||
type BudgetHandler struct{}
|
||||
|
||||
func NewBudgetHandler() *BudgetHandler { return &BudgetHandler{} }
|
||||
|
||||
// budgetResponse is the canonical JSON shape for both GET and PATCH responses.
|
||||
type budgetResponse struct {
|
||||
// BudgetLimit is the monthly spend ceiling in USD cents (null = no limit).
|
||||
// budget_limit=500 means $5.00/month.
|
||||
BudgetLimit *int64 `json:"budget_limit"`
|
||||
// MonthlySpend is the agent's self-reported accumulated LLM API spend
|
||||
// for the current month (USD cents). Incremented via heartbeat.
|
||||
MonthlySpend int64 `json:"monthly_spend"`
|
||||
// BudgetRemaining is null when BudgetLimit is null, otherwise
|
||||
// max(0, budget_limit - monthly_spend). Can be negative — we store the
|
||||
// actual value so callers can see how far over-budget a workspace is.
|
||||
BudgetRemaining *int64 `json:"budget_remaining"`
|
||||
}
|
||||
|
||||
// GetBudget handles GET /workspaces/:id/budget.
|
||||
// Returns the workspace's current budget ceiling, accumulated spend, and
|
||||
// computed remaining headroom. Both budget_limit and budget_remaining are
|
||||
// null when no limit has been configured for the workspace.
|
||||
func (h *BudgetHandler) GetBudget(c *gin.Context) {
|
||||
workspaceID := c.Param("id")
|
||||
ctx := c.Request.Context()
|
||||
|
||||
var budgetLimit sql.NullInt64
|
||||
var monthlySpend int64
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT budget_limit, COALESCE(monthly_spend, 0)
|
||||
FROM workspaces
|
||||
WHERE id = $1 AND status != 'removed'`,
|
||||
workspaceID,
|
||||
).Scan(&budgetLimit, &monthlySpend)
|
||||
if err == sql.ErrNoRows {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"})
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
log.Printf("GetBudget: query failed for %s: %v", workspaceID, err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "query failed"})
|
||||
return
|
||||
}
|
||||
|
||||
resp := budgetResponse{
|
||||
MonthlySpend: monthlySpend,
|
||||
}
|
||||
if budgetLimit.Valid {
|
||||
limit := budgetLimit.Int64
|
||||
resp.BudgetLimit = &limit
|
||||
remaining := limit - monthlySpend
|
||||
resp.BudgetRemaining = &remaining
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// patchBudgetRequest is the expected JSON body for PATCH /workspaces/:id/budget.
|
||||
// budget_limit=null removes the ceiling; a positive integer sets it (USD cents).
|
||||
type patchBudgetRequest struct {
|
||||
// BudgetLimit pointer so JSON null → nil, absent → parse error (required field).
|
||||
BudgetLimit *int64 `json:"budget_limit"`
|
||||
}
|
||||
|
||||
// PatchBudget handles PATCH /workspaces/:id/budget.
|
||||
// Accepts {"budget_limit": <int64>} to set a new ceiling, or
|
||||
// {"budget_limit": null} to remove an existing ceiling.
|
||||
// Returns the updated budget state in the same shape as GetBudget.
|
||||
func (h *BudgetHandler) PatchBudget(c *gin.Context) {
|
||||
workspaceID := c.Param("id")
|
||||
ctx := c.Request.Context()
|
||||
|
||||
// We need to distinguish between "field absent" and "field = null",
|
||||
// so we unmarshal into a raw map first.
|
||||
var raw map[string]interface{}
|
||||
if err := c.ShouldBindJSON(&raw); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
budgetLimitRaw, ok := raw["budget_limit"]
|
||||
if !ok {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit field is required"})
|
||||
return
|
||||
}
|
||||
|
||||
// Validate and convert the value. JSON numbers decode as float64.
|
||||
var budgetArg interface{} // nil → SQL NULL, int64 → new ceiling
|
||||
if budgetLimitRaw != nil {
|
||||
switch v := budgetLimitRaw.(type) {
|
||||
case float64:
|
||||
if v < 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit must be >= 0 (USD cents)"})
|
||||
return
|
||||
}
|
||||
cv := int64(v)
|
||||
budgetArg = cv
|
||||
case int64:
|
||||
if v < 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit must be >= 0 (USD cents)"})
|
||||
return
|
||||
}
|
||||
budgetArg = v
|
||||
default:
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "budget_limit must be an integer (USD cents) or null"})
|
||||
return
|
||||
}
|
||||
}
|
||||
// budgetArg == nil means "clear the ceiling"
|
||||
|
||||
// Existence check — return 404 for non-existent / removed workspaces.
|
||||
var exists bool
|
||||
if err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT EXISTS(SELECT 1 FROM workspaces WHERE id = $1 AND status != 'removed')`,
|
||||
workspaceID,
|
||||
).Scan(&exists); err != nil || !exists {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "workspace not found"})
|
||||
return
|
||||
}
|
||||
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE workspaces SET budget_limit = $2, updated_at = now() WHERE id = $1`,
|
||||
workspaceID, budgetArg,
|
||||
); err != nil {
|
||||
log.Printf("PatchBudget: update failed for %s: %v", workspaceID, err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "update failed"})
|
||||
return
|
||||
}
|
||||
|
||||
// Re-read the current state so the response reflects exactly what is in
|
||||
// the DB, including the monthly_spend the agent has already accumulated.
|
||||
var newLimit sql.NullInt64
|
||||
var monthlySpend int64
|
||||
if err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT budget_limit, COALESCE(monthly_spend, 0) FROM workspaces WHERE id = $1`,
|
||||
workspaceID,
|
||||
).Scan(&newLimit, &monthlySpend); err != nil {
|
||||
log.Printf("PatchBudget: re-read failed for %s: %v", workspaceID, err)
|
||||
// Still success — just omit the echo.
|
||||
c.JSON(http.StatusOK, gin.H{"status": "updated"})
|
||||
return
|
||||
}
|
||||
|
||||
resp := budgetResponse{
|
||||
MonthlySpend: monthlySpend,
|
||||
}
|
||||
if newLimit.Valid {
|
||||
limit := newLimit.Int64
|
||||
resp.BudgetLimit = &limit
|
||||
remaining := limit - monthlySpend
|
||||
resp.BudgetRemaining = &remaining
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, resp)
|
||||
}
|
||||
458
platform/internal/handlers/budget_test.go
Normal file
458
platform/internal/handlers/budget_test.go
Normal file
@ -0,0 +1,458 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// ==================== GET /workspaces/:id/budget ====================
|
||||
|
||||
// TestBudgetGet_NotFound verifies that GET /budget returns 404 for an unknown
|
||||
// workspace ID (ErrNoRows from the budget query).
|
||||
func TestBudgetGet_NotFound(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\)`).
|
||||
WithArgs("ws-not-there").
|
||||
WillReturnError(sql.ErrNoRows)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-not-there"}}
|
||||
c.Request = httptest.NewRequest("GET", "/workspaces/ws-not-there/budget", nil)
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.GetBudget(c)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Errorf("expected 404, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetGet_DBError verifies that a non-ErrNoRows DB error returns 500.
|
||||
func TestBudgetGet_DBError(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\)`).
|
||||
WithArgs("ws-db-err").
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-db-err"}}
|
||||
c.Request = httptest.NewRequest("GET", "/workspaces/ws-db-err/budget", nil)
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.GetBudget(c)
|
||||
|
||||
if w.Code != http.StatusInternalServerError {
|
||||
t.Errorf("expected 500, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetGet_NoLimit verifies that budget_limit and budget_remaining are
|
||||
// null when the workspace has no budget ceiling configured.
|
||||
func TestBudgetGet_NoLimit(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\)`).
|
||||
WithArgs("ws-free").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}).
|
||||
AddRow(nil, int64(42)))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-free"}}
|
||||
c.Request = httptest.NewRequest("GET", "/workspaces/ws-free/budget", nil)
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.GetBudget(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("parse response: %v", err)
|
||||
}
|
||||
if resp["budget_limit"] != nil {
|
||||
t.Errorf("expected budget_limit=null, got %v", resp["budget_limit"])
|
||||
}
|
||||
if resp["budget_remaining"] != nil {
|
||||
t.Errorf("expected budget_remaining=null, got %v", resp["budget_remaining"])
|
||||
}
|
||||
if resp["monthly_spend"] != float64(42) {
|
||||
t.Errorf("expected monthly_spend=42, got %v", resp["monthly_spend"])
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetGet_WithLimit verifies that budget_limit, monthly_spend, and
|
||||
// budget_remaining are all returned correctly when a ceiling is set.
|
||||
func TestBudgetGet_WithLimit(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\)`).
|
||||
WithArgs("ws-capped").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}).
|
||||
AddRow(int64(500), int64(123)))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-capped"}}
|
||||
c.Request = httptest.NewRequest("GET", "/workspaces/ws-capped/budget", nil)
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.GetBudget(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("parse response: %v", err)
|
||||
}
|
||||
if resp["budget_limit"] != float64(500) {
|
||||
t.Errorf("expected budget_limit=500, got %v", resp["budget_limit"])
|
||||
}
|
||||
if resp["monthly_spend"] != float64(123) {
|
||||
t.Errorf("expected monthly_spend=123, got %v", resp["monthly_spend"])
|
||||
}
|
||||
// budget_remaining = 500 - 123 = 377
|
||||
if resp["budget_remaining"] != float64(377) {
|
||||
t.Errorf("expected budget_remaining=377, got %v", resp["budget_remaining"])
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetGet_OverBudget verifies that budget_remaining can be negative
|
||||
// when monthly_spend has already exceeded budget_limit.
|
||||
func TestBudgetGet_OverBudget(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\)`).
|
||||
WithArgs("ws-over").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}).
|
||||
AddRow(int64(100), int64(150)))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-over"}}
|
||||
c.Request = httptest.NewRequest("GET", "/workspaces/ws-over/budget", nil)
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.GetBudget(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("parse response: %v", err)
|
||||
}
|
||||
// budget_remaining = 100 - 150 = -50 (negative, but we store actual value)
|
||||
if resp["budget_remaining"] != float64(-50) {
|
||||
t.Errorf("expected budget_remaining=-50, got %v", resp["budget_remaining"])
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== PATCH /workspaces/:id/budget ====================
|
||||
|
||||
// TestBudgetPatch_MissingField verifies that PATCH /budget with no budget_limit
|
||||
// field in the body returns 400.
|
||||
func TestBudgetPatch_MissingField(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-patch-missing"}}
|
||||
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-patch-missing/budget",
|
||||
bytes.NewBufferString(`{"other_field":123}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.PatchBudget(c)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetPatch_InvalidBody verifies that a malformed JSON body returns 400.
|
||||
func TestBudgetPatch_InvalidBody(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-patch-bad"}}
|
||||
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-patch-bad/budget",
|
||||
bytes.NewBufferString(`not json`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.PatchBudget(c)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetPatch_NegativeValue verifies that a negative budget_limit is rejected.
|
||||
func TestBudgetPatch_NegativeValue(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-negative"}}
|
||||
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-negative/budget",
|
||||
bytes.NewBufferString(`{"budget_limit":-1}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.PatchBudget(c)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400 for negative budget_limit, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetPatch_InvalidType verifies that a non-numeric budget_limit returns 400.
|
||||
func TestBudgetPatch_InvalidType(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-badtype"}}
|
||||
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-badtype/budget",
|
||||
bytes.NewBufferString(`{"budget_limit":"not-a-number"}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.PatchBudget(c)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("expected 400 for string budget_limit, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetPatch_WorkspaceNotFound verifies that PATCH /budget returns 404
|
||||
// when the workspace doesn't exist.
|
||||
func TestBudgetPatch_WorkspaceNotFound(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`).
|
||||
WithArgs("ws-no-exist").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-no-exist"}}
|
||||
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-no-exist/budget",
|
||||
bytes.NewBufferString(`{"budget_limit":500}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.PatchBudget(c)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Errorf("expected 404, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetPatch_SetLimit verifies that PATCH /budget with a positive value
|
||||
// updates the DB and returns the new budget state.
|
||||
func TestBudgetPatch_SetLimit(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
// Existence probe
|
||||
mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`).
|
||||
WithArgs("ws-set-limit").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
|
||||
// UPDATE
|
||||
mock.ExpectExec(`UPDATE workspaces SET budget_limit`).
|
||||
WithArgs("ws-set-limit", int64(500)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
// Re-read for response
|
||||
mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\) FROM workspaces WHERE id`).
|
||||
WithArgs("ws-set-limit").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}).
|
||||
AddRow(int64(500), int64(200)))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-set-limit"}}
|
||||
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-set-limit/budget",
|
||||
bytes.NewBufferString(`{"budget_limit":500}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.PatchBudget(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("parse response: %v", err)
|
||||
}
|
||||
if resp["budget_limit"] != float64(500) {
|
||||
t.Errorf("expected budget_limit=500, got %v", resp["budget_limit"])
|
||||
}
|
||||
if resp["monthly_spend"] != float64(200) {
|
||||
t.Errorf("expected monthly_spend=200, got %v", resp["monthly_spend"])
|
||||
}
|
||||
// budget_remaining = 500 - 200 = 300
|
||||
if resp["budget_remaining"] != float64(300) {
|
||||
t.Errorf("expected budget_remaining=300, got %v", resp["budget_remaining"])
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetPatch_ClearLimit verifies that PATCH /budget with budget_limit=null
|
||||
// clears the ceiling, making budget_limit and budget_remaining null in the response.
|
||||
func TestBudgetPatch_ClearLimit(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`).
|
||||
WithArgs("ws-clear-limit").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
|
||||
// UPDATE with NULL
|
||||
mock.ExpectExec(`UPDATE workspaces SET budget_limit`).
|
||||
WithArgs("ws-clear-limit", nil).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
// Re-read — budget_limit is now NULL
|
||||
mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\) FROM workspaces WHERE id`).
|
||||
WithArgs("ws-clear-limit").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}).
|
||||
AddRow(nil, int64(50)))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-clear-limit"}}
|
||||
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-clear-limit/budget",
|
||||
bytes.NewBufferString(`{"budget_limit":null}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.PatchBudget(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
||||
t.Fatalf("parse response: %v", err)
|
||||
}
|
||||
if resp["budget_limit"] != nil {
|
||||
t.Errorf("expected budget_limit=null after clear, got %v", resp["budget_limit"])
|
||||
}
|
||||
if resp["budget_remaining"] != nil {
|
||||
t.Errorf("expected budget_remaining=null after clear, got %v", resp["budget_remaining"])
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetPatch_UpdateDBError verifies that a DB error during the UPDATE
|
||||
// returns 500.
|
||||
func TestBudgetPatch_UpdateDBError(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`).
|
||||
WithArgs("ws-patch-dberr").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
|
||||
mock.ExpectExec(`UPDATE workspaces SET budget_limit`).
|
||||
WithArgs("ws-patch-dberr", int64(500)).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-patch-dberr"}}
|
||||
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-patch-dberr/budget",
|
||||
bytes.NewBufferString(`{"budget_limit":500}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.PatchBudget(c)
|
||||
|
||||
if w.Code != http.StatusInternalServerError {
|
||||
t.Errorf("expected 500 on UPDATE error, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBudgetPatch_ZeroLimit verifies that budget_limit=0 is accepted (it means
|
||||
// every A2A call is blocked — useful to pause a workspace's LLM spend entirely).
|
||||
func TestBudgetPatch_ZeroLimit(t *testing.T) {
|
||||
mock := setupTestDB(t)
|
||||
setupTestRedis(t)
|
||||
|
||||
mock.ExpectQuery(`SELECT EXISTS.*status != 'removed'`).
|
||||
WithArgs("ws-zero-limit").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
|
||||
mock.ExpectExec(`UPDATE workspaces SET budget_limit`).
|
||||
WithArgs("ws-zero-limit", int64(0)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectQuery(`SELECT budget_limit, COALESCE\(monthly_spend, 0\) FROM workspaces WHERE id`).
|
||||
WithArgs("ws-zero-limit").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"budget_limit", "monthly_spend"}).
|
||||
AddRow(int64(0), int64(0)))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Params = gin.Params{{Key: "id", Value: "ws-zero-limit"}}
|
||||
c.Request = httptest.NewRequest("PATCH", "/workspaces/ws-zero-limit/budget",
|
||||
bytes.NewBufferString(`{"budget_limit":0}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
h := NewBudgetHandler()
|
||||
h.PatchBudget(c)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("expected 200 for zero budget_limit, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Errorf("sqlmock expectations not met: %v", err)
|
||||
}
|
||||
}
|
||||
@ -256,6 +256,12 @@ func Setup(hub *ws.Hub, broadcaster *events.Broadcaster, prov *provisioner.Provi
|
||||
// (mirrors the /workspaces/:id/a2a pattern). Issue #249.
|
||||
r.GET("/workspaces/:id/schedules/health", schedh.Health)
|
||||
|
||||
// Budget — per-workspace spend ceiling and current usage (#541).
|
||||
// GET returns the current state; PATCH sets or clears the ceiling.
|
||||
budgeth := handlers.NewBudgetHandler()
|
||||
wsAuth.GET("/budget", budgeth.GetBudget)
|
||||
wsAuth.PATCH("/budget", budgeth.PatchBudget)
|
||||
|
||||
// Token management (user-facing create/list/revoke)
|
||||
tokh := handlers.NewTokenHandler()
|
||||
wsAuth.GET("/tokens", tokh.List)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user