feat(management): add usage record retrieval endpoint

- Implemented `/v0/management/usage` endpoint for fetching queued usage records from Redis.
- Included validation for `count` parameter to ensure positive integers.
- Added unit tests for queue retrieval and validation, with authentication validation in integration tests.
- Updated management routing to include the new endpoint.
This commit is contained in:
Luis Pater
2026-05-05 02:53:04 +08:00
parent ba5d8ca733
commit 61b39d49bd
4 changed files with 209 additions and 0 deletions
+55
View File
@@ -0,0 +1,55 @@
package management
import (
"encoding/json"
"errors"
"net/http"
"strconv"
"strings"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue"
)
type usageQueueRecord []byte
func (r usageQueueRecord) MarshalJSON() ([]byte, error) {
if json.Valid(r) {
return append([]byte(nil), r...), nil
}
return json.Marshal(string(r))
}
// GetUsage pops queued usage records from the Redis-compatible usage queue.
func (h *Handler) GetUsage(c *gin.Context) {
if h == nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "handler unavailable"})
return
}
count, errCount := parseUsageQueueCount(c.Query("count"))
if errCount != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": errCount.Error()})
return
}
items := redisqueue.PopOldest(count)
records := make([]usageQueueRecord, 0, len(items))
for _, item := range items {
records = append(records, usageQueueRecord(append([]byte(nil), item...)))
}
c.JSON(http.StatusOK, records)
}
func parseUsageQueueCount(value string) (int, error) {
value = strings.TrimSpace(value)
if value == "" {
return 1, nil
}
count, errCount := strconv.Atoi(value)
if errCount != nil || count <= 0 {
return 0, errors.New("count must be a positive integer")
}
return count, nil
}
@@ -0,0 +1,98 @@
package management
import (
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
"github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue"
)
func TestGetUsagePopsRequestedRecords(t *testing.T) {
gin.SetMode(gin.TestMode)
withManagementUsageQueue(t, func() {
redisqueue.Enqueue([]byte(`{"id":1}`))
redisqueue.Enqueue([]byte(`{"id":2}`))
redisqueue.Enqueue([]byte(`{"id":3}`))
rec := httptest.NewRecorder()
ginCtx, _ := gin.CreateTestContext(rec)
ginCtx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/usage?count=2", nil)
h := &Handler{}
h.GetUsage(ginCtx)
if rec.Code != http.StatusOK {
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusOK, rec.Body.String())
}
var payload []json.RawMessage
if errUnmarshal := json.Unmarshal(rec.Body.Bytes(), &payload); errUnmarshal != nil {
t.Fatalf("unmarshal response: %v", errUnmarshal)
}
if len(payload) != 2 {
t.Fatalf("response records = %d, want 2", len(payload))
}
requireRecordID(t, payload[0], 1)
requireRecordID(t, payload[1], 2)
remaining := redisqueue.PopOldest(10)
if len(remaining) != 1 || string(remaining[0]) != `{"id":3}` {
t.Fatalf("remaining queue = %q, want third item only", remaining)
}
})
}
func TestGetUsageInvalidCountDoesNotPop(t *testing.T) {
gin.SetMode(gin.TestMode)
withManagementUsageQueue(t, func() {
redisqueue.Enqueue([]byte(`{"id":1}`))
rec := httptest.NewRecorder()
ginCtx, _ := gin.CreateTestContext(rec)
ginCtx.Request = httptest.NewRequest(http.MethodGet, "/v0/management/usage?count=0", nil)
h := &Handler{}
h.GetUsage(ginCtx)
if rec.Code != http.StatusBadRequest {
t.Fatalf("status = %d, want %d body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
}
remaining := redisqueue.PopOldest(10)
if len(remaining) != 1 || string(remaining[0]) != `{"id":1}` {
t.Fatalf("remaining queue = %q, want original item", remaining)
}
})
}
func withManagementUsageQueue(t *testing.T, fn func()) {
t.Helper()
prevQueueEnabled := redisqueue.Enabled()
redisqueue.SetEnabled(false)
redisqueue.SetEnabled(true)
defer func() {
redisqueue.SetEnabled(false)
redisqueue.SetEnabled(prevQueueEnabled)
}()
fn()
}
func requireRecordID(t *testing.T, raw json.RawMessage, want int) {
t.Helper()
var payload struct {
ID int `json:"id"`
}
if errUnmarshal := json.Unmarshal(raw, &payload); errUnmarshal != nil {
t.Fatalf("unmarshal record: %v", errUnmarshal)
}
if payload.ID != want {
t.Fatalf("record id = %d, want %d", payload.ID, want)
}
}