270 lines
11 KiB
Go
270 lines
11 KiB
Go
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 TestServiceCloseSubscriptionManagedKeyOverridesExplicitProbeAPIKey(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: "host-user-1", APIKey: "managed-user-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 host.gatewayProbe.APIKey != "managed-user-key" {
|
|
t.Fatalf("gateway probe api key = %q, want managed-user-key override", host.gatewayProbe.APIKey)
|
|
}
|
|
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 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
|
|
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) {
|
|
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
|
|
}
|