package access import ( "context" "errors" "testing" "sub2api-cn-relay-manager/internal/host/sub2api" ) func TestValidateRejectsMissingProbeAPIKeyForSelfService(t *testing.T) { err := Validate(ClosureRequest{Mode: "self_service"}) if err == nil { t.Fatal("Validate() error = nil, want validation error") } } func TestValidateRejectsMissingSubscriptionsForSubscriptionMode(t *testing.T) { err := Validate(ClosureRequest{Mode: "subscription", ProbeAPIKey: "user-key"}) if err == nil { t.Fatal("Validate() error = nil, want validation error") } } func TestValidateAllowsManagedSubscriptionProbeWithoutExplicitAPIKey(t *testing.T) { err := Validate(ClosureRequest{ Mode: "subscription", GroupID: "group-1", ExpectedModel: "deepseek-chat", Subscriptions: []SubscriptionTarget{{UserID: "crm-user-42", DurationDays: 30}}, }) if err != nil { t.Fatalf("Validate() error = %v, want nil for managed subscription probe", err) } } func TestServiceCloseAssignsSubscriptionsAndProbesGateway(t *testing.T) { host := &fakeClosureHost{ gatewayResult: sub2api.GatewayAccessResult{OK: true, StatusCode: 200, HasExpectedModel: true, Models: []string{"deepseek-chat"}}, completionResult: sub2api.GatewayCompletionResult{OK: true, StatusCode: 200, ContentType: "application/json"}, managedAccess: map[string]sub2api.SubscriptionAccessRef{ "user-1": {UserID: "host-user-1", APIKey: "managed-user-key"}, }, } service := NewService(host) result, err := service.Close(context.Background(), ClosureRequest{ Mode: "subscription", GroupID: "group-1", ExpectedModel: "deepseek-chat", Subscriptions: []SubscriptionTarget{{UserID: "user-1", DurationDays: 30}}, }) if err != nil { t.Fatalf("Close() error = %v", err) } if len(host.assigned) != 1 { t.Fatalf("assigned subscriptions = %d, want 1", len(host.assigned)) } if host.assigned[0].UserID != "host-user-1" { t.Fatalf("assigned subscription user = %q, want host-user-1", host.assigned[0].UserID) } if host.gatewayProbe.APIKey != "managed-user-key" || host.gatewayProbe.ExpectedModel != "deepseek-chat" { t.Fatalf("gateway probe = %+v, want api key + expected model", host.gatewayProbe) } if host.completionProbe.APIKey != "managed-user-key" || host.completionProbe.Model != "deepseek-chat" { t.Fatalf("completion probe = %+v, want api key + model", host.completionProbe) } if !result.OK || !result.HasExpectedModel { t.Fatalf("gateway result = %+v, want success", result) } if !result.CompletionOK { t.Fatalf("completion result = %+v, want success", result) } if result.EffectiveProbeAPIKey != "managed-user-key" { t.Fatalf("effective probe api key = %q, want managed-user-key", result.EffectiveProbeAPIKey) } if result.EffectiveProbeKeySource != ProbeKeySourceManagedSubscription { t.Fatalf("effective probe key source = %q, want %q", result.EffectiveProbeKeySource, ProbeKeySourceManagedSubscription) } } func TestServiceCloseSubscriptionExplicitProbeAPIKeyUsesRequestedUserKey(t *testing.T) { host := &fakeClosureHost{ gatewayResult: sub2api.GatewayAccessResult{OK: true, StatusCode: 200, HasExpectedModel: true, Models: []string{"deepseek-chat"}}, completionResult: sub2api.GatewayCompletionResult{OK: true, StatusCode: 200}, managedAccess: map[string]sub2api.SubscriptionAccessRef{ "user-1": {UserID: "user-1", APIKey: "caller-supplied-key"}, }, } service := NewService(host) result, err := service.Close(context.Background(), ClosureRequest{ Mode: "subscription", ProbeAPIKey: "caller-supplied-key", GroupID: "group-1", ExpectedModel: "deepseek-chat", Subscriptions: []SubscriptionTarget{{UserID: "user-1", DurationDays: 30}}, }) if err != nil { t.Fatalf("Close() error = %v", err) } if len(host.assigned) != 1 || host.assigned[0].UserID != "user-1" { t.Fatalf("assigned subscriptions = %+v, want requested user assignment", host.assigned) } if host.gatewayProbe.APIKey != "caller-supplied-key" { t.Fatalf("gateway probe api key = %q, want caller-supplied-key", host.gatewayProbe.APIKey) } if len(host.ensureRequests) != 1 || host.ensureRequests[0].ProbeAPIKey != "caller-supplied-key" { t.Fatalf("ensure requests = %+v, want probe api key forwarded", host.ensureRequests) } if result.EffectiveProbeAPIKey != "caller-supplied-key" { t.Fatalf("effective probe api key = %q, want caller-supplied-key", result.EffectiveProbeAPIKey) } if result.EffectiveProbeKeySource != ProbeKeySourceRequestedProbeAPIKey { t.Fatalf("effective probe key source = %q, want %q", result.EffectiveProbeKeySource, ProbeKeySourceRequestedProbeAPIKey) } } func TestServiceCloseReturnsSubscriptionErrorBeforeGatewayProbe(t *testing.T) { host := &fakeClosureHost{assignErr: errors.New("assign failed")} service := NewService(host) _, err := service.Close(context.Background(), ClosureRequest{ Mode: "subscription", ProbeAPIKey: "user-key", GroupID: "group-1", ExpectedModel: "deepseek-chat", Subscriptions: []SubscriptionTarget{{UserID: "user-1", DurationDays: 30}}, }) if err == nil { t.Fatal("Close() error = nil, want subscription failure") } if host.gatewayProbe.APIKey != "" { t.Fatalf("gateway probe should not run after subscription error, got %+v", host.gatewayProbe) } } func TestServiceCloseRetriesTransientGatewayCompletionFailure(t *testing.T) { host := &fakeClosureHost{ gatewayResult: sub2api.GatewayAccessResult{OK: true, StatusCode: 200, HasExpectedModel: true, Models: []string{"kimi-k2.6"}}, completionResults: []sub2api.GatewayCompletionResult{ {OK: false, StatusCode: 503, ContentType: "application/json", BodyPreview: `{"error":{"message":"Service temporarily unavailable","type":"api_error"}}`}, {OK: true, StatusCode: 200, ContentType: "application/json"}, }, managedAccess: map[string]sub2api.SubscriptionAccessRef{ "user-1": {UserID: "host-user-1", APIKey: "managed-user-key"}, }, } result, err := NewService(host).Close(context.Background(), ClosureRequest{ Mode: "subscription", GroupID: "group-1", ExpectedModel: "kimi-k2.6", Subscriptions: []SubscriptionTarget{{UserID: "user-1", DurationDays: 30}}, }) if err != nil { t.Fatalf("Close() error = %v", err) } if !result.CompletionOK || result.CompletionStatus != 200 { t.Fatalf("completion result = %+v, want retried success", result) } if host.completionCalls != 2 { t.Fatalf("completion calls = %d, want 2", host.completionCalls) } } func TestServiceCloseRepairsOpenAIResponsesCapabilityMismatch(t *testing.T) { host := &fakeClosureHost{ gatewayResult: sub2api.GatewayAccessResult{OK: true, StatusCode: 200, HasExpectedModel: true, Models: []string{"kimi-k2.6"}}, completionResults: []sub2api.GatewayCompletionResult{ {OK: false, StatusCode: 502, ContentType: "application/json", BodyPreview: `{"error":{"message":"Upstream service temporarily unavailable","type":"upstream_error"}}`}, }, completionAfterRepair: &sub2api.GatewayCompletionResult{OK: true, StatusCode: 200, ContentType: "application/json"}, managedAccess: map[string]sub2api.SubscriptionAccessRef{ "user-1": {UserID: "host-user-1", APIKey: "managed-user-key"}, }, } result, err := NewService(host).Close(context.Background(), ClosureRequest{ Mode: "subscription", GroupID: "group-1", AccountIDs: []string{"account-1", "account-1"}, ExpectedModel: "kimi-k2.6", ResponsesCapabilitySuspect: true, Subscriptions: []SubscriptionTarget{{UserID: "user-1", DurationDays: 30}}, }) if err != nil { t.Fatalf("Close() error = %v", err) } if !result.CompletionOK || result.CompletionStatus != 200 { t.Fatalf("completion result = %+v, want repaired success", result) } if host.disableResponsesCalls != 1 { t.Fatalf("disable responses calls = %d, want 1", host.disableResponsesCalls) } if len(host.disabledResponsesAccountIDs) != 1 || host.disabledResponsesAccountIDs[0] != "account-1" { t.Fatalf("disabled responses account ids = %v, want [account-1]", host.disabledResponsesAccountIDs) } if host.clearTempUnschedulableCalls != 1 { t.Fatalf("clear temp unschedulable calls = %d, want 1", host.clearTempUnschedulableCalls) } if len(host.clearedTempUnschedulableAccountIDs) != 1 || host.clearedTempUnschedulableAccountIDs[0] != "account-1" { t.Fatalf("cleared temp unschedulable account ids = %v, want [account-1]", host.clearedTempUnschedulableAccountIDs) } } type fakeClosureHost struct { assigned []sub2api.AssignSubscriptionRequest ensureRequests []sub2api.EnsureSubscriptionAccessRequest managedAccess map[string]sub2api.SubscriptionAccessRef assignErr error gatewayProbe sub2api.GatewayAccessCheckRequest gatewayResult sub2api.GatewayAccessResult gatewayErr error completionProbe sub2api.GatewayCompletionCheckRequest completionCalls int completionResults []sub2api.GatewayCompletionResult completionResult sub2api.GatewayCompletionResult completionAfterRepair *sub2api.GatewayCompletionResult completionErr error disableResponsesCalls int disabledResponsesAccountIDs []string clearTempUnschedulableCalls int clearedTempUnschedulableAccountIDs []string } func (f *fakeClosureHost) EnsureSubscriptionAccess(_ context.Context, req sub2api.EnsureSubscriptionAccessRequest) (sub2api.SubscriptionAccessRef, error) { f.ensureRequests = append(f.ensureRequests, req) if ref, ok := f.managedAccess[req.UserSelector]; ok { return ref, nil } return sub2api.SubscriptionAccessRef{}, errors.New("missing managed access") } func (f *fakeClosureHost) AssignSubscription(_ context.Context, req sub2api.AssignSubscriptionRequest) (sub2api.SubscriptionRef, error) { if f.assignErr != nil { return sub2api.SubscriptionRef{}, f.assignErr } f.assigned = append(f.assigned, req) return sub2api.SubscriptionRef{ID: "sub-1"}, nil } func (f *fakeClosureHost) CheckGatewayAccess(_ context.Context, req sub2api.GatewayAccessCheckRequest) (sub2api.GatewayAccessResult, error) { f.gatewayProbe = req if f.gatewayErr != nil { return sub2api.GatewayAccessResult{}, f.gatewayErr } return f.gatewayResult, nil } func (f *fakeClosureHost) CheckGatewayCompletion(_ context.Context, req sub2api.GatewayCompletionCheckRequest) (sub2api.GatewayCompletionResult, error) { f.completionProbe = req f.completionCalls++ if f.completionErr != nil { return sub2api.GatewayCompletionResult{}, f.completionErr } if f.disableResponsesCalls > 0 && f.completionAfterRepair != nil { return *f.completionAfterRepair, nil } if len(f.completionResults) > 0 { idx := f.completionCalls - 1 if idx >= len(f.completionResults) { idx = len(f.completionResults) - 1 } return f.completionResults[idx], nil } return f.completionResult, nil } func (f *fakeClosureHost) DisableOpenAIResponsesAPI(_ context.Context, accountIDs []string) error { f.disableResponsesCalls++ f.disabledResponsesAccountIDs = append([]string(nil), accountIDs...) return nil } func (f *fakeClosureHost) ClearTempUnschedulable(_ context.Context, accountIDs []string) error { f.clearTempUnschedulableCalls++ f.clearedTempUnschedulableAccountIDs = append([]string(nil), accountIDs...) return nil }