internal/registry/model_registry_transient_suspend_deadline_test.go
Rama main · réplica de lectura del repositorio original.
package registry
import (
"testing"
"time"
)
func listedIDs(models []map[string]any) map[string]bool {
out := make(map[string]bool, len(models))
for _, m := range models {
if id, _ := m["id"].(string); id != "" {
out[id] = true
}
}
return out
}
func infoIDs(infos []*ModelInfo) map[string]bool {
out := make(map[string]bool, len(infos))
for _, info := range infos {
if info != nil && info.ID != "" {
out[info.ID] = true
}
}
return out
}
func TestModelRegistry_TransientSuspendDeadline_AutoExpiresCacheAndAllViews(t *testing.T) {
baseNow := time.Date(2026, 11, 11, 1, 5, 0, 0, time.UTC)
mockNow := baseNow
r := newTestModelRegistry()
r.nowFunc = func() time.Time { return mockNow }
r.RegisterClient("xai-auth", "xai", []*ModelInfo{
{ID: "grok-4.7", OwnedBy: "xai", Created: baseNow.Unix()},
{ID: "grok-4.7-cp", OwnedBy: "xai", Created: baseNow.Unix()},
{ID: "grok-4.6", OwnedBy: "xai", Created: baseNow.Unix()},
})
epoch := r.ClientRegistrationEpoch("xai-auth")
deadline := baseNow.Add(time.Minute)
if !r.ApplyClientModelProjections("xai-auth", epoch, 1, []ClientModelProjection{
{ModelID: "grok-4.7", Suspended: true, SuspendReason: "xai stream error: stream disconnected", SuspendUntil: deadline},
{ModelID: "grok-4.7-cp", Suspended: true, SuspendReason: "xai stream error: stream disconnected", SuspendUntil: deadline},
{ModelID: "grok-4.6"},
}) {
t.Fatal("expected initial projection to apply")
}
// During cooldown: grok-4.7 and grok-4.7-cp are hidden, grok-4.6 remains listed.
duringOpenAI := listedIDs(r.GetAvailableModels("openai"))
if duringOpenAI["grok-4.7"] || duringOpenAI["grok-4.7-cp"] || !duringOpenAI["grok-4.6"] {
t.Fatalf("unexpected /v1/models set during transient cooldown: %v", duringOpenAI)
}
if cached := r.availableModelsCache["openai"]; !cached.expiresAt.Equal(deadline) {
t.Fatalf("availableModelsCache expiresAt = %v, want %v", cached.expiresAt, deadline)
}
if infos := infoIDs(r.GetAvailableModelInfos()); infos["grok-4.7"] || infos["grok-4.7-cp"] || !infos["grok-4.6"] {
t.Fatalf("unexpected GetAvailableModelInfos during transient cooldown: %v", infos)
}
if provInfos := infoIDs(r.GetAvailableModelsByProvider("xai")); provInfos["grok-4.7"] || provInfos["grok-4.7-cp"] || !provInfos["grok-4.6"] {
t.Fatalf("unexpected GetAvailableModelsByProvider during transient cooldown: %v", provInfos)
}
if got := r.GetModelCount("grok-4.7"); got != 0 {
t.Fatalf("GetModelCount(grok-4.7) during cooldown = %d, want 0", got)
}
if !r.IsModelSuspendedForClient("xai-auth", "grok-4.7") {
t.Fatal("expected IsModelSuspendedForClient to report true during active transient cooldown")
}
// Still active 1ns before deadline.
mockNow = deadline.Add(-time.Nanosecond)
stillCooling := listedIDs(r.GetAvailableModels("openai"))
if stillCooling["grok-4.7"] || stillCooling["grok-4.7-cp"] {
t.Fatalf("models listed before deadline elapsed: %v", stillCooling)
}
// Exact deadline instant with NO re-projection, NO request, and NO refresh.
mockNow = deadline
afterOpenAI := listedIDs(r.GetAvailableModels("openai"))
if !afterOpenAI["grok-4.7"] || !afterOpenAI["grok-4.7-cp"] || !afterOpenAI["grok-4.6"] {
t.Fatalf("expected all models listed once transient cooldown expired without re-projection, got %v", afterOpenAI)
}
if infos := infoIDs(r.GetAvailableModelInfos()); !infos["grok-4.7"] || !infos["grok-4.7-cp"] {
t.Fatalf("expected GetAvailableModelInfos to list models after deadline, got %v", infos)
}
if provInfos := infoIDs(r.GetAvailableModelsByProvider("xai")); !provInfos["grok-4.7"] || !provInfos["grok-4.7-cp"] {
t.Fatalf("expected GetAvailableModelsByProvider to list models after deadline, got %v", provInfos)
}
if got := r.GetModelCount("grok-4.7"); got != 1 {
t.Fatalf("GetModelCount(grok-4.7) after cooldown = %d, want 1", got)
}
if got := r.GetModelCount("grok-4.7-cp"); got != 1 {
t.Fatalf("GetModelCount(grok-4.7-cp) after cooldown = %d, want 1", got)
}
if r.IsModelSuspendedForClient("xai-auth", "grok-4.7") {
t.Fatal("expected IsModelSuspendedForClient to report false once transient cooldown expired")
}
r.CleanupExpiredQuotas()
if _, exists := r.models["grok-4.7"].SuspendedClients["xai-auth"]; exists {
t.Fatal("expected CleanupExpiredQuotas to prune expired transient suspension entry")
}
if _, exists := r.models["grok-4.7"].SuspendedClientDeadlines["xai-auth"]; exists {
t.Fatal("expected CleanupExpiredQuotas to prune expired transient suspension deadline")
}
}
func TestModelRegistry_OpenEndedSuspensionAndManualOverridePersist(t *testing.T) {
baseNow := time.Date(2026, 11, 11, 2, 0, 0, 0, time.UTC)
mockNow := baseNow
r := newTestModelRegistry()
r.nowFunc = func() time.Time { return mockNow }
r.RegisterClient("xai-auth", "xai", []*ModelInfo{{ID: "grok-4.7", OwnedBy: "xai"}})
epoch := r.ClientRegistrationEpoch("xai-auth")
deadline := baseNow.Add(time.Minute)
// Start with a bounded transient suspension.
if !r.ApplyClientModelProjections("xai-auth", epoch, 1, []ClientModelProjection{
{ModelID: "grok-4.7", Suspended: true, SuspendReason: "transient 408", SuspendUntil: deadline},
}) {
t.Fatal("initial transient projection rejected")
}
// Override with manual open-ended SuspendClientModel while already in SuspendedClients.
r.SuspendClientModel("xai-auth", "grok-4.7", "manual")
mockNow = baseNow.Add(24 * time.Hour)
if listed := listedIDs(r.GetAvailableModels("openai")); listed["grok-4.7"] {
t.Fatalf("manual SuspendClientModel leaked model after earlier transient deadline: %v", listed)
}
if !r.IsModelSuspendedForClient("xai-auth", "grok-4.7") {
t.Fatal("expected manual open-ended suspension to stay active")
}
// Resume, then apply an open-ended projection (zero SuspendUntil).
r.ResumeClientModel("xai-auth", "grok-4.7")
if !listedIDs(r.GetAvailableModels("openai"))["grok-4.7"] {
t.Fatal("expected ResumeClientModel to restore listing")
}
if !r.ApplyClientModelProjections("xai-auth", epoch, 2, []ClientModelProjection{
{ModelID: "grok-4.7", Suspended: true, SuspendReason: "disabled"},
}) {
t.Fatal("open-ended projection rejected")
}
mockNow = baseNow.Add(48 * time.Hour)
if listed := listedIDs(r.GetAvailableModels("openai")); listed["grok-4.7"] {
t.Fatalf("open-ended projection leaked model: %v", listed)
}
}
func TestModelRegistry_ShortQuotaDeadlineRecoversWithDisabledBackup(t *testing.T) {
baseNow := time.Date(2026, 11, 11, 3, 0, 0, 0, time.UTC)
mockNow := baseNow
r := newTestModelRegistry()
r.nowFunc = func() time.Time { return mockNow }
model := []*ModelInfo{{ID: "grok-4.7", OwnedBy: "xai"}}
r.RegisterClient("c1-quota", "xai", model)
r.RegisterClient("c2-disabled", "xai", model)
epoch1 := r.ClientRegistrationEpoch("c1-quota")
epoch2 := r.ClientRegistrationEpoch("c2-disabled")
shortQuotaDeadline := baseNow.Add(time.Minute)
if !r.ApplyClientModelProjections("c1-quota", epoch1, 1, []ClientModelProjection{
{ModelID: "grok-4.7", Suspended: true, SuspendReason: "quota", SuspendUntil: shortQuotaDeadline, QuotaExceeded: true},
}) {
t.Fatal("quota projection rejected")
}
if !r.ApplyClientModelProjections("c2-disabled", epoch2, 1, []ClientModelProjection{
{ModelID: "grok-4.7", Suspended: true, SuspendReason: "disabled"},
}) {
t.Fatal("disabled projection rejected")
}
if listedIDs(r.GetAvailableModels("openai"))["grok-4.7"] {
t.Fatal("expected model hidden while c1 is in 1m quota cooldown and c2 is disabled")
}
// Advance past the 1m quota deadline (still within the default 5m modelQuotaExceededWindow).
mockNow = shortQuotaDeadline.Add(time.Second)
if !listedIDs(r.GetAvailableModels("openai"))["grok-4.7"] {
t.Fatal("expected model listed once c1's 1m quota deadline passed")
}
if got := r.GetModelCount("grok-4.7"); got != 1 {
t.Fatalf("GetModelCount(grok-4.7) after c1 quota recovery = %d, want 1", got)
}
}