202 lines
6.6 KiB
Go
202 lines
6.6 KiB
Go
package plugin
|
|
|
|
import (
|
|
"context"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"path/filepath"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
managedaccess "cpa-ext/internal/access"
|
|
"cpa-ext/internal/repository"
|
|
)
|
|
|
|
func lifecycleRequest(t *testing.T, schema uint32, config string) []byte {
|
|
t.Helper()
|
|
raw, err := json.Marshal(LifecycleRequest{ConfigYAML: []byte(config), SchemaVersion: schema})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return raw
|
|
}
|
|
|
|
func testConfig(t *testing.T, config string) string {
|
|
t.Helper()
|
|
return fmt.Sprintf("database_path: %q\n%s", filepath.Join(t.TempDir(), "usage.db"), config)
|
|
}
|
|
|
|
func TestRegisterNegotiatesSchemaAndDeclaresManagedAccessCapabilities(t *testing.T) {
|
|
app := NewApp()
|
|
raw, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, 99, testConfig(t, "enabled: true\ncodex_only: true\n")))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var env Envelope
|
|
if err := json.Unmarshal(raw, &env); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if !env.OK {
|
|
t.Fatalf("register failed: %s", raw)
|
|
}
|
|
var got Registration
|
|
if err := json.Unmarshal(env.Result, &got); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if got.SchemaVersion != SchemaVersion || !got.Capabilities.FrontendAuthProvider || !got.Capabilities.FrontendAuthProviderExclusive ||
|
|
!got.Capabilities.Scheduler || !got.Capabilities.RequestInterceptor || !got.Capabilities.RequestLifecyclePlugin ||
|
|
!got.Capabilities.UsagePlugin || !got.Capabilities.ManagementAPI {
|
|
t.Fatalf("unexpected registration: %+v", got)
|
|
}
|
|
}
|
|
|
|
func TestLifecycleConfigYAMLUsesBase64WireEncoding(t *testing.T) {
|
|
raw := lifecycleRequest(t, SchemaVersion, "enabled: true\n")
|
|
if !json.Valid(raw) {
|
|
t.Fatal("invalid JSON")
|
|
}
|
|
var wire map[string]any
|
|
_ = json.Unmarshal(raw, &wire)
|
|
want := base64.StdEncoding.EncodeToString([]byte("enabled: true\n"))
|
|
if wire["config_yaml"] != want {
|
|
t.Fatalf("config_yaml = %v, want %q", wire["config_yaml"], want)
|
|
}
|
|
}
|
|
|
|
func TestUsageRecordWireFieldsMatchTargetContract(t *testing.T) {
|
|
raw, err := json.Marshal(UsageRecord{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var fields map[string]json.RawMessage
|
|
if err := json.Unmarshal(raw, &fields); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
want := []string{
|
|
"Provider", "ExecutorType", "Model", "Alias", "APIKey", "AuthID", "AuthIndex", "AuthType",
|
|
"Source", "ReasoningEffort", "ServiceTier", "Generate", "RequestedAt", "Latency", "TTFT",
|
|
"Failed", "Failure", "Detail", "ResponseHeaders",
|
|
}
|
|
if len(fields) != len(want) {
|
|
t.Fatalf("UsageRecord fields=%v", fields)
|
|
}
|
|
for _, name := range want {
|
|
if _, ok := fields[name]; !ok {
|
|
t.Fatalf("UsageRecord missing target field %s", name)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestReconfigureIsAtomicAndFiltersUsage(t *testing.T) {
|
|
app := NewApp()
|
|
if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: true\n"))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for _, record := range []UsageRecord{{Provider: "codex", Model: "gpt-5.5"}, {Provider: "gemini", Model: "gemini-pro"}} {
|
|
raw, _ := json.Marshal(record)
|
|
if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
if got := app.Seen(); got != 1 {
|
|
t.Fatalf("seen = %d, want 1", got)
|
|
}
|
|
if _, err := app.HandleMethod(MethodPluginReconfigure, lifecycleRequest(t, SchemaVersion, "enabled: [invalid")); err == nil {
|
|
t.Fatal("invalid reconfiguration unexpectedly succeeded")
|
|
}
|
|
raw, _ := json.Marshal(UsageRecord{Provider: "codex", Model: "gpt-5.5"})
|
|
_, _ = app.HandleMethod(MethodUsageHandle, raw)
|
|
if got := app.Seen(); got != 2 {
|
|
t.Fatalf("last valid config was not retained: seen = %d", got)
|
|
}
|
|
}
|
|
|
|
func TestReconfigureReleasesAdmissionsWhenDatabaseChanges(t *testing.T) {
|
|
directory := t.TempDir()
|
|
firstPath := filepath.Join(directory, "first.db")
|
|
secondPath := filepath.Join(directory, "second.db")
|
|
store, err := repository.OpenSQLiteUsage(secondPath)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := store.BootstrapManagedKey(context.Background(), "default", "000000"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := store.UpdateBilling(context.Background(), "key_default", managedaccess.BillingSettings{
|
|
QuotaMicros: 1_000_000, ResetPeriod: managedaccess.ResetNone, MaxConcurrency: 4,
|
|
}, time.Now()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := store.AuthorizeBilling(context.Background(), "key_default", "stale-request", time.Now()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.Close(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
app := NewApp()
|
|
if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, fmt.Sprintf("database_path: %q\n", firstPath))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := app.store.UpdateBilling(context.Background(), "key_default", managedaccess.BillingSettings{
|
|
QuotaMicros: 1_000_000, ResetPeriod: managedaccess.ResetNone, MaxConcurrency: 4,
|
|
}, time.Now()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := app.store.AuthorizeBilling(context.Background(), "key_default", "live-request", time.Now()); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := app.HandleMethod(MethodPluginReconfigure, lifecycleRequest(t, SchemaVersion, fmt.Sprintf("database_path: %q\n", firstPath))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
state, err := app.store.BillingState(context.Background(), "key_default", time.Now())
|
|
if err != nil || state.ActiveRequests != 1 {
|
|
t.Fatalf("same-database reconfiguration released live admission: state=%+v err=%v", state, err)
|
|
}
|
|
if _, err := app.HandleMethod(MethodPluginReconfigure, lifecycleRequest(t, SchemaVersion, fmt.Sprintf("database_path: %q\n", secondPath))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
state, err = app.store.BillingState(context.Background(), "key_default", time.Now())
|
|
if err != nil || state.ActiveRequests != 0 {
|
|
t.Fatalf("reconfigured billing state=%+v err=%v", state, err)
|
|
}
|
|
}
|
|
|
|
func TestConcurrentUsage(t *testing.T) {
|
|
app := NewApp()
|
|
if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: true\n"))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
raw, _ := json.Marshal(UsageRecord{Provider: "codex", Model: "gpt-5.5"})
|
|
var wg sync.WaitGroup
|
|
for range 100 {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil {
|
|
t.Errorf("handle usage: %v", err)
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
if got := app.Seen(); got != 100 {
|
|
t.Fatalf("seen = %d, want 100", got)
|
|
}
|
|
}
|
|
|
|
func TestUnknownMethodReturnsErrorEnvelope(t *testing.T) {
|
|
raw, err := NewApp().HandleMethod("missing", nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var env Envelope
|
|
if err := json.Unmarshal(raw, &env); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if env.OK || env.Error == nil || env.Error.Code != "unknown_method" {
|
|
t.Fatalf("unexpected envelope: %s", raw)
|
|
}
|
|
}
|