mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-05-22 04:14:45 +08:00
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:
55
internal/api/handlers/management/usage.go
Normal file
55
internal/api/handlers/management/usage.go
Normal 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
|
||||
}
|
||||
98
internal/api/handlers/management/usage_test.go
Normal file
98
internal/api/handlers/management/usage_test.go
Normal file
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -551,6 +551,7 @@ func (s *Server) registerManagementRoutes() {
|
||||
mgmt.PATCH("/api-keys", s.mgmt.PatchAPIKeys)
|
||||
mgmt.DELETE("/api-keys", s.mgmt.DeleteAPIKeys)
|
||||
mgmt.GET("/api-key-usage", s.mgmt.GetAPIKeyUsage)
|
||||
mgmt.GET("/usage", s.mgmt.GetUsage)
|
||||
|
||||
mgmt.GET("/gemini-api-key", s.mgmt.GetGeminiKeys)
|
||||
mgmt.PUT("/gemini-api-key", s.mgmt.PutGeminiKeys)
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
gin "github.com/gin-gonic/gin"
|
||||
proxyconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/redisqueue"
|
||||
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
|
||||
@@ -84,6 +85,60 @@ func TestHealthz(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestManagementUsageRequiresManagementAuthAndPopsArray(t *testing.T) {
|
||||
t.Setenv("MANAGEMENT_PASSWORD", "test-management-key")
|
||||
|
||||
prevQueueEnabled := redisqueue.Enabled()
|
||||
redisqueue.SetEnabled(false)
|
||||
t.Cleanup(func() {
|
||||
redisqueue.SetEnabled(false)
|
||||
redisqueue.SetEnabled(prevQueueEnabled)
|
||||
})
|
||||
|
||||
server := newTestServer(t)
|
||||
|
||||
redisqueue.Enqueue([]byte(`{"id":1}`))
|
||||
redisqueue.Enqueue([]byte(`{"id":2}`))
|
||||
|
||||
missingKeyReq := httptest.NewRequest(http.MethodGet, "/v0/management/usage?count=2", nil)
|
||||
missingKeyRR := httptest.NewRecorder()
|
||||
server.engine.ServeHTTP(missingKeyRR, missingKeyReq)
|
||||
if missingKeyRR.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("missing key status = %d, want %d body=%s", missingKeyRR.Code, http.StatusUnauthorized, missingKeyRR.Body.String())
|
||||
}
|
||||
|
||||
authReq := httptest.NewRequest(http.MethodGet, "/v0/management/usage?count=2", nil)
|
||||
authReq.Header.Set("Authorization", "Bearer test-management-key")
|
||||
authRR := httptest.NewRecorder()
|
||||
server.engine.ServeHTTP(authRR, authReq)
|
||||
if authRR.Code != http.StatusOK {
|
||||
t.Fatalf("authenticated status = %d, want %d body=%s", authRR.Code, http.StatusOK, authRR.Body.String())
|
||||
}
|
||||
|
||||
var payload []json.RawMessage
|
||||
if errUnmarshal := json.Unmarshal(authRR.Body.Bytes(), &payload); errUnmarshal != nil {
|
||||
t.Fatalf("unmarshal response: %v body=%s", errUnmarshal, authRR.Body.String())
|
||||
}
|
||||
if len(payload) != 2 {
|
||||
t.Fatalf("response records = %d, want 2", len(payload))
|
||||
}
|
||||
for i, raw := range payload {
|
||||
var record struct {
|
||||
ID int `json:"id"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(raw, &record); errUnmarshal != nil {
|
||||
t.Fatalf("unmarshal record %d: %v", i, errUnmarshal)
|
||||
}
|
||||
if record.ID != i+1 {
|
||||
t.Fatalf("record %d id = %d, want %d", i, record.ID, i+1)
|
||||
}
|
||||
}
|
||||
|
||||
if remaining := redisqueue.PopOldest(1); len(remaining) != 0 {
|
||||
t.Fatalf("remaining queue = %q, want empty", remaining)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAmpProviderModelRoutes(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
|
||||
Reference in New Issue
Block a user