sub2api-cn-relay-manager/internal/access/closure_test.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)
	}
}

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)
	_, 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)
	}
}

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)
	}
}

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
	completionErr     error
}

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 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
}