sub2api-cn-relay-manager/internal/provision/pool_routing_test.go

package provision

import (
	"context"
	"strings"
	"testing"

	"sub2api-cn-relay-manager/internal/host/sub2api"
	"sub2api-cn-relay-manager/internal/pack"
)

func TestPoolRoutingWithDualVendors(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	store := openProvisionTestStore(t)
	defer closeProvisionTestStore(t, store)

	seedProvisionHost(t, store, "host-a", "https://api-a.example.com")
	seedProvisionHost(t, store, "host-b", "https://api-b.example.com")

	providerA := pack.ProviderManifest{
		ProviderID:      "deepseek-official",
		DisplayName:     "DeepSeek Official",
		BaseURL:         "https://api.deepseek.com",
		Platform:        "openai",
		AccountType:     "apikey",
		DefaultModels:   []string{"deepseek-chat", "deepseek-reasoner"},
		SmokeTestModel:  "deepseek-chat",
		GroupTemplate:   pack.GroupTemplate{Name: "DeepSeek Official Group", RateMultiplier: 1},
		ChannelTemplate: pack.ChannelTemplate{Name: "DeepSeek Official Channel", ModelMapping: map[string]string{"deepseek-chat": "deepseek-chat"}},
		PlanTemplate:    pack.PlanTemplate{Name: "DeepSeek Plan", Price: 19.9, ValidityDays: 30, ValidityUnit: "day"},
	}

	providerB := pack.ProviderManifest{
		ProviderID:      "deepseek-backup",
		DisplayName:     "DeepSeek Backup Proxy",
		BaseURL:         "https://backup.deepseek.example.com",
		Platform:        "openai",
		AccountType:     "apikey",
		DefaultModels:   []string{"deepseek-chat", "deepseek-reasoner"},
		SmokeTestModel:  "deepseek-chat",
		GroupTemplate:   pack.GroupTemplate{Name: "DeepSeek Backup Group", RateMultiplier: 1},
		ChannelTemplate: pack.ChannelTemplate{Name: "DeepSeek Backup Channel", ModelMapping: map[string]string{"deepseek-chat": "deepseek-chat"}},
		PlanTemplate:    pack.PlanTemplate{Name: "DeepSeek Backup Plan", Price: 19.9, ValidityDays: 30, ValidityUnit: "day"},
	}

	packManifest := pack.LoadedPack{
		Manifest: pack.Manifest{
			PackID:         "openai-cn-pack",
			Version:        "1.0.0",
			TargetHost:     "sub2api",
			MinHostVersion: "0.1.126",
			MaxHostVersion: "0.2.x",
		},
		Checksum: "checksum-1",
	}

	// Import provider A via host-a
	hostA := &fakeHostAdapter{
		batchAccounts: []sub2api.AccountRef{{ID: "account_a1"}},
		testResults:   map[string]sub2api.ProbeResult{"account_a1": {OK: true, Status: "passed"}},
		models:        map[string][]sub2api.AccountModel{"account_a1": {{ID: "deepseek-chat"}, {ID: "deepseek-reasoner"}}},
		gatewayResult: sub2api.GatewayAccessResult{OK: true, StatusCode: 200, HasExpectedModel: true, Models: []string{"deepseek-chat"}, CompletionOK: true, CompletionStatus: 200},
	}

	resultA, errA := NewRuntimeImportService(store, hostA).Import(ctx, RuntimeImportRequest{
		HostID:      "host-a",
		HostBaseURL: "https://api-a.example.com",
		Pack:        packManifest,
		Provider:    providerA,
		Mode:        ImportModePartial,
		Keys:        []string{"sk-key-a"},
		Access:      AccessRequest{Mode: AccessModeSelfService, ProbeAPIKey: "user-key"},
	})
	if errA != nil {
		t.Fatalf("deepseek-official Import() error = %v", errA)
	}
	if resultA.BatchID <= 0 {
		t.Fatalf("BatchID = %d for provider A, want positive", resultA.BatchID)
	}

	// Import provider B via host-b
	hostB := &fakeHostAdapter{
		batchAccounts: []sub2api.AccountRef{{ID: "account_b1"}},
		testResults:   map[string]sub2api.ProbeResult{"account_b1": {OK: true, Status: "passed"}},
		models:        map[string][]sub2api.AccountModel{"account_b1": {{ID: "deepseek-chat"}, {ID: "deepseek-reasoner"}}},
		gatewayResult: sub2api.GatewayAccessResult{OK: true, StatusCode: 200, HasExpectedModel: true, Models: []string{"deepseek-chat"}, CompletionOK: true, CompletionStatus: 200},
	}

	resultB, errB := NewRuntimeImportService(store, hostB).Import(ctx, RuntimeImportRequest{
		HostID:      "host-b",
		HostBaseURL: "https://api-b.example.com",
		Pack:        packManifest,
		Provider:    providerB,
		Mode:        ImportModePartial,
		Keys:        []string{"sk-key-b"},
		Access:      AccessRequest{Mode: AccessModeSelfService, ProbeAPIKey: "user-key"},
	})
	if errB != nil {
		t.Fatalf("deepseek-backup Import() error = %v", errB)
	}
	if resultB.BatchID <= 0 {
		t.Fatalf("BatchID = %d for provider B, want positive", resultB.BatchID)
	}

	// Verify each provider has its own logical_group_model with deepseek-chat
	groupsA, err := store.LogicalGroupModels().ListByLogicalGroupID(ctx, resultA.Report.Group.ID)
	if err != nil {
		t.Fatalf("ListByLogicalGroupID(A) error = %v", err)
	}
	if len(groupsA) != 1 || groupsA[0].PublicModel != "deepseek-chat" {
		t.Fatalf("group A models = %+v, want 1 model [deepseek-chat]", groupsA)
	}

	groupsB, err := store.LogicalGroupModels().ListByLogicalGroupID(ctx, resultB.Report.Group.ID)
	if err != nil {
		t.Fatalf("ListByLogicalGroupID(B) error = %v", err)
	}
	if len(groupsB) != 1 || groupsB[0].PublicModel != "deepseek-chat" {
		t.Fatalf("group B models = %+v, want 1 model [deepseek-chat]", groupsB)
	}

	// Verify each provider has its own logical_group_route
	routesA, err := store.LogicalGroupRoutes().ListByLogicalGroupID(ctx, resultA.Report.Group.ID)
	if err != nil {
		t.Fatalf("ListByLogicalGroupID routes A error = %v", err)
	}
	if len(routesA) != 1 {
		t.Fatalf("routes for group A = %d, want 1", len(routesA))
	}
	if !strings.HasPrefix(routesA[0].RouteID, "route-") {
		t.Fatalf("route A RouteID = %q, want route-* prefix", routesA[0].RouteID)
	}

	routesB, err := store.LogicalGroupRoutes().ListByLogicalGroupID(ctx, resultB.Report.Group.ID)
	if err != nil {
		t.Fatalf("ListByLogicalGroupID routes B error = %v", err)
	}
	if len(routesB) != 1 {
		t.Fatalf("routes for group B = %d, want 1", len(routesB))
	}

	// Verify each route carries deepseek-chat route model
	rmA, err := store.LogicalGroupRouteModels().ListByRouteID(ctx, routesA[0].RouteID)
	if err != nil {
		t.Fatalf("ListByRouteID models A error = %v", err)
	}
	hasChatA := false
	for _, rm := range rmA {
		if rm.PublicModel == "deepseek-chat" {
			hasChatA = true
			break
		}
	}
	if !hasChatA {
		t.Fatalf("route A models = %+v, want deepseek-chat", rmA)
	}

	rmB, err := store.LogicalGroupRouteModels().ListByRouteID(ctx, routesB[0].RouteID)
	if err != nil {
		t.Fatalf("ListByRouteID models B error = %v", err)
	}
	hasChatB := false
	for _, rm := range rmB {
		if rm.PublicModel == "deepseek-chat" {
			hasChatB = true
			break
		}
	}
	if !hasChatB {
		t.Fatalf("route B models = %+v, want deepseek-chat", rmB)
	}
}