mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-06-09 15:53:18 +08:00
feat(api): implement protocol multiplexer and Redis queue for usage integration
- Added `protocol_multiplexer.go`, enabling support for both HTTP and Redis protocols on a single listener. - Introduced `redis_queue_protocol.go` to handle Redis-compatible RESP commands for queue management. - Integrated `redisqueue` package, supporting in-memory queuing with expiration pruning. - Updated server initialization to manage a shared listener and multiplex connections. - Adjusted `Handler` to adopt `AuthenticateManagementKey` for modular key validation, supporting both HTTP and Redis flows.
This commit is contained in:
160
internal/redisqueue/plugin_test.go
Normal file
160
internal/redisqueue/plugin_test.go
Normal file
@@ -0,0 +1,160 @@
|
||||
package redisqueue
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
internallogging "github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
||||
internalusage "github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/usage"
|
||||
)
|
||||
|
||||
func TestUsageQueuePluginPayloadIncludesStableFieldsAndSuccess(t *testing.T) {
|
||||
withEnabledQueue(t, func() {
|
||||
ginCtx := newTestGinContext(t, http.MethodPost, "/v1/chat/completions", http.StatusOK)
|
||||
internallogging.SetGinRequestID(ginCtx, "gin-request-id-ignored")
|
||||
ctx := context.WithValue(internallogging.WithRequestID(context.Background(), "ctx-request-id"), "gin", ginCtx)
|
||||
|
||||
plugin := &usageQueuePlugin{}
|
||||
plugin.HandleUsage(ctx, coreusage.Record{
|
||||
Provider: "openai",
|
||||
Model: "gpt-5.4",
|
||||
APIKey: "test-key",
|
||||
AuthIndex: "0",
|
||||
AuthType: "apikey",
|
||||
Source: "user@example.com",
|
||||
RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC),
|
||||
Latency: 1500 * time.Millisecond,
|
||||
Detail: coreusage.Detail{
|
||||
InputTokens: 10,
|
||||
OutputTokens: 20,
|
||||
TotalTokens: 30,
|
||||
},
|
||||
})
|
||||
|
||||
payload := popSinglePayload(t)
|
||||
requireStringField(t, payload, "provider", "openai")
|
||||
requireStringField(t, payload, "model", "gpt-5.4")
|
||||
requireStringField(t, payload, "endpoint", "POST /v1/chat/completions")
|
||||
requireStringField(t, payload, "auth_type", "apikey")
|
||||
requireStringField(t, payload, "request_id", "ctx-request-id")
|
||||
requireBoolField(t, payload, "failed", false)
|
||||
})
|
||||
}
|
||||
|
||||
func TestUsageQueuePluginPayloadIncludesStableFieldsAndFailureAndGinRequestID(t *testing.T) {
|
||||
withEnabledQueue(t, func() {
|
||||
ginCtx := newTestGinContext(t, http.MethodGet, "/v1/responses", http.StatusInternalServerError)
|
||||
internallogging.SetGinRequestID(ginCtx, "gin-request-id")
|
||||
ctx := context.WithValue(context.Background(), "gin", ginCtx)
|
||||
|
||||
plugin := &usageQueuePlugin{}
|
||||
plugin.HandleUsage(ctx, coreusage.Record{
|
||||
Provider: "openai",
|
||||
Model: "gpt-5.4-mini",
|
||||
APIKey: "test-key",
|
||||
AuthIndex: "0",
|
||||
AuthType: "apikey",
|
||||
Source: "user@example.com",
|
||||
RequestedAt: time.Date(2026, 4, 25, 0, 0, 0, 0, time.UTC),
|
||||
Latency: 2500 * time.Millisecond,
|
||||
Detail: coreusage.Detail{
|
||||
InputTokens: 10,
|
||||
OutputTokens: 20,
|
||||
TotalTokens: 30,
|
||||
},
|
||||
})
|
||||
|
||||
payload := popSinglePayload(t)
|
||||
requireStringField(t, payload, "provider", "openai")
|
||||
requireStringField(t, payload, "model", "gpt-5.4-mini")
|
||||
requireStringField(t, payload, "endpoint", "GET /v1/responses")
|
||||
requireStringField(t, payload, "auth_type", "apikey")
|
||||
requireStringField(t, payload, "request_id", "gin-request-id")
|
||||
requireBoolField(t, payload, "failed", true)
|
||||
})
|
||||
}
|
||||
|
||||
func withEnabledQueue(t *testing.T, fn func()) {
|
||||
t.Helper()
|
||||
|
||||
prevQueueEnabled := Enabled()
|
||||
prevStatsEnabled := internalusage.StatisticsEnabled()
|
||||
|
||||
SetEnabled(false)
|
||||
SetEnabled(true)
|
||||
internalusage.SetStatisticsEnabled(true)
|
||||
|
||||
defer func() {
|
||||
SetEnabled(false)
|
||||
SetEnabled(prevQueueEnabled)
|
||||
internalusage.SetStatisticsEnabled(prevStatsEnabled)
|
||||
}()
|
||||
|
||||
fn()
|
||||
}
|
||||
|
||||
func newTestGinContext(t *testing.T, method, path string, status int) *gin.Context {
|
||||
t.Helper()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
ginCtx, _ := gin.CreateTestContext(recorder)
|
||||
ginCtx.Request = httptest.NewRequest(method, "http://example.com"+path, nil)
|
||||
if status != 0 {
|
||||
ginCtx.Status(status)
|
||||
}
|
||||
return ginCtx
|
||||
}
|
||||
|
||||
func popSinglePayload(t *testing.T) map[string]json.RawMessage {
|
||||
t.Helper()
|
||||
|
||||
items := PopOldest(10)
|
||||
if len(items) != 1 {
|
||||
t.Fatalf("PopOldest() items = %d, want 1", len(items))
|
||||
}
|
||||
|
||||
var payload map[string]json.RawMessage
|
||||
if err := json.Unmarshal(items[0], &payload); err != nil {
|
||||
t.Fatalf("unmarshal payload: %v", err)
|
||||
}
|
||||
return payload
|
||||
}
|
||||
|
||||
func requireStringField(t *testing.T, payload map[string]json.RawMessage, key, want string) {
|
||||
t.Helper()
|
||||
|
||||
raw, ok := payload[key]
|
||||
if !ok {
|
||||
t.Fatalf("payload missing %q", key)
|
||||
}
|
||||
var got string
|
||||
if err := json.Unmarshal(raw, &got); err != nil {
|
||||
t.Fatalf("unmarshal %q: %v", key, err)
|
||||
}
|
||||
if got != want {
|
||||
t.Fatalf("%s = %q, want %q", key, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func requireBoolField(t *testing.T, payload map[string]json.RawMessage, key string, want bool) {
|
||||
t.Helper()
|
||||
|
||||
raw, ok := payload[key]
|
||||
if !ok {
|
||||
t.Fatalf("payload missing %q", key)
|
||||
}
|
||||
var got bool
|
||||
if err := json.Unmarshal(raw, &got); err != nil {
|
||||
t.Fatalf("unmarshal %q: %v", key, err)
|
||||
}
|
||||
if got != want {
|
||||
t.Fatalf("%s = %t, want %t", key, got, want)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user