niuniu/sub2api-cn-relay-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(batch): add live reuse admin verification flow

Browse Source
This commit is contained in:
phamnazage-jpg
2026-05-27 20:23:42 +08:00
parent 02580cda0b
commit ebd86a4256
16 changed files with 1768 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

272
internal/app/http_batch_import_test.go
View File

@@ -2,12 +2,16 @@ package app
import (
"context"
"crypto/sha256"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"sub2api-cn-relay-manager/internal/batch"
"sub2api-cn-relay-manager/internal/pack"
"sub2api-cn-relay-manager/internal/store/sqlite"
"sub2api-cn-relay-manager/internal/testutil"
)
@@ -197,6 +201,120 @@ func TestBatchImportHTTP(t *testing.T) {
t.Fatalf("item = %+v, want current_stage=done and access_status=active", items[0])
}
})
t.Run("create run action reuses matched legacy account", func(t *testing.T) {
t.Parallel()
server := httptest.NewServer(newBatchImportActionStubServer(t))
defer server.Close()
dsn := testutil.SQLiteTestDSN(t, "reuse-state.db", true)
store := testutil.OpenSQLiteStore(t, dsn)
defer closeAppTestStore(t, store)
hostPK := mustCreateBatchImportActionHost(t, store, server.URL)
packPK, providerPK := mustSeedLegacyBatchImportProvider(t, store, server.URL)
legacyBatchID := mustCreateLegacyReusableBatch(t, store, hostPK, packPK, providerPK, "entry-key", "account_1")
action := buildCreateBatchImportRunAction(dsn)
result, err := action(context.Background(), CreateBatchImportRunRequest{
HostID: "host-1",
Mode: "strict",
AccessMode: "self_service",
ConfirmWaitTimeoutSec: 1,
ProbeAPIKey: "gateway-key",
Entries: []BatchImportEntryRequest{{
BaseURL: server.URL,
APIKey: "entry-key",
RequestedModels: []string{"kimi-k2.6"},
}},
})
if err != nil {
t.Fatalf("buildCreateBatchImportRunAction() reuse error = %v", err)
}
if strings.TrimSpace(result.RunID) == "" {
t.Fatalf("result.RunID = %q, want non-empty", result.RunID)
}
items, err := store.ImportRunItems().ListByRunID(context.Background(), result.RunID)
if err != nil {
t.Fatalf("ImportRunItems().ListByRunID() reuse error = %v", err)
}
if len(items) != 1 {
t.Fatalf("len(items) = %d, want 1", len(items))
}
item := items[0]
if !item.ProvisionReused || item.AccountResolution != "reused" || item.MatchedAccountState != "active" {
t.Fatalf("reuse item = %+v, want provision_reused + reused + active", item)
}
if item.LegacyBatchID == nil || *item.LegacyBatchID != legacyBatchID {
t.Fatalf("LegacyBatchID = %v, want %d", item.LegacyBatchID, legacyBatchID)
}
if item.CurrentStage != "done" || item.AccessStatus != "active" {
t.Fatalf("reuse item final state = %+v, want done/active", item)
}
batches, err := store.ImportBatches().ListByProviderIDAndHostID(context.Background(), providerPK, hostPK)
if err != nil {
t.Fatalf("ImportBatches().ListByProviderIDAndHostID() error = %v", err)
}
if len(batches) != 1 {
t.Fatalf("len(batches) = %d, want 1 legacy batch only", len(batches))
}
})
t.Run("create run action reuses legacy account when pack provider id differs from normalized runtime id", func(t *testing.T) {
t.Parallel()
server := httptest.NewServer(newBatchImportActionStubServer(t))
defer server.Close()
dsn := testutil.SQLiteTestDSN(t, "reuse-baseurl-fallback.db", true)
store := testutil.OpenSQLiteStore(t, dsn)
defer closeAppTestStore(t, store)
hostPK := mustCreateBatchImportActionHost(t, store, server.URL)
packPK, providerPK := mustSeedLegacyBatchImportProviderWithID(t, store, server.URL, "legacy-pack-provider")
legacyBatchID := mustCreateLegacyReusableBatch(t, store, hostPK, packPK, providerPK, "entry-key", "101")
action := buildCreateBatchImportRunAction(dsn)
result, err := action(context.Background(), CreateBatchImportRunRequest{
HostID: "host-1",
Mode: "strict",
AccessMode: "self_service",
ConfirmWaitTimeoutSec: 1,
ProbeAPIKey: "gateway-key",
Entries: []BatchImportEntryRequest{{
BaseURL: server.URL,
APIKey: "entry-key",
RequestedModels: []string{"kimi-k2.6"},
}},
})
if err != nil {
t.Fatalf("buildCreateBatchImportRunAction() base_url fallback reuse error = %v", err)
}
if strings.TrimSpace(result.RunID) == "" {
t.Fatalf("result.RunID = %q, want non-empty", result.RunID)
}
items, err := store.ImportRunItems().ListByRunID(context.Background(), result.RunID)
if err != nil {
t.Fatalf("ImportRunItems().ListByRunID() base_url fallback error = %v", err)
}
if len(items) != 1 {
t.Fatalf("len(items) = %d, want 1", len(items))
}
item := items[0]
if !item.ProvisionReused || item.AccountResolution != "reused" || item.MatchedAccountState != "active" {
t.Fatalf("reuse item = %+v, want provision_reused + reused + active", item)
}
if item.LegacyBatchID == nil || *item.LegacyBatchID != legacyBatchID {
t.Fatalf("LegacyBatchID = %v, want %d", item.LegacyBatchID, legacyBatchID)
}
if item.CurrentStage != "done" || item.AccessStatus != "active" {
t.Fatalf("reuse item final state = %+v, want done/active", item)
}
})
}
func TestBatchImportWrapperFunctions(t *testing.T) {
@@ -363,6 +481,21 @@ func newBatchImportActionStubServer(t *testing.T) http.Handler {
}
writeJSON(w, http.StatusOK, map[string]any{"data": map[string]any{"items": []map[string]any{{"id": "kimi-k2.6", "display_name": "Kimi K2.6", "type": "chat"}}}})
})
mux.HandleFunc("/api/v1/admin/accounts/101/test", func(w http.ResponseWriter, r *http.Request) {
if !requireBatchImportActionAdminToken(t, w, r) {
return
}
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("event: result\n"))
_, _ = w.Write([]byte("data: {\"status\":\"passed\",\"message\":\"smoke passed\",\"ok\":true}\n\n"))
})
mux.HandleFunc("/api/v1/admin/accounts/101/models", func(w http.ResponseWriter, r *http.Request) {
if !requireBatchImportActionAdminToken(t, w, r) {
return
}
writeJSON(w, http.StatusOK, map[string]any{"data": map[string]any{"items": []map[string]any{{"id": "kimi-k2.6", "display_name": "Kimi K2.6", "type": "chat"}}}})
})
mux.HandleFunc("/api/v1/admin/subscriptions/assign", func(w http.ResponseWriter, r *http.Request) {
if !requireBatchImportActionAdminToken(t, w, r) {
return
@@ -425,3 +558,142 @@ func requireBatchImportActionAdminToken(t *testing.T, w http.ResponseWriter, r *
}
return true
}
func mustCreateBatchImportActionHost(t *testing.T, store *sqlite.DB, baseURL string) int64 {
t.Helper()
hostPK, err := store.Hosts().Create(context.Background(), sqlite.Host{
HostID: "host-1",
BaseURL: baseURL,
HostVersion: "0.1.126",
CapabilityProbeJSON: "{}",
AuthType: "apikey",
AuthToken: "host-token",
})
if err != nil {
t.Fatalf("Hosts().Create() error = %v", err)
}
return hostPK
}
func mustSeedLegacyBatchImportProvider(t *testing.T, store *sqlite.DB, baseURL string) (int64, int64) {
t.Helper()
return mustSeedLegacyBatchImportProviderWithID(t, store, baseURL, batch.NormalizeProviderID(baseURL))
}
func mustSeedLegacyBatchImportProviderWithID(t *testing.T, store *sqlite.DB, baseURL, providerID string) (int64, int64) {
t.Helper()
packPK, err := store.Packs().Create(context.Background(), sqlite.Pack{
PackID: "seed-pack",
Version: "1.0.0",
Checksum: "seed-pack@1.0.0",
Vendor: "test",
TargetHost: "sub2api",
ManifestJSON: `{"pack_id":"seed-pack","version":"1.0.0","target_host":"sub2api"}`,
})
if err != nil {
t.Fatalf("Packs().Create() error = %v", err)
}
providerManifest := pack.ProviderManifest{
ProviderID: strings.TrimSpace(providerID),
DisplayName: "Legacy Reuse Provider",
BaseURL: baseURL,
Platform: "openai",
AccountType: "apikey",
DefaultModels: []string{"kimi-k2.6"},
SmokeTestModel: "kimi-k2.6",
GroupTemplate: pack.GroupTemplate{
Name: "legacy-group",
RateMultiplier: 1,
},
ChannelTemplate: pack.ChannelTemplate{
Name: "legacy-channel",
ModelMapping: map[string]string{"kimi-k2.6": "kimi-k2.6"},
},
PlanTemplate: pack.PlanTemplate{
Name: "legacy-plan",
Price: 1,
ValidityDays: 30,
ValidityUnit: "day",
},
Import: pack.ImportOptions{
SupportsMultiKey: true,
SupportsStrict: true,
SupportsPartial: true,
},
}
manifestJSON, err := json.Marshal(providerManifest)
if err != nil {
t.Fatalf("json.Marshal(providerManifest) error = %v", err)
}
defaultModelsJSON, err := json.Marshal(providerManifest.DefaultModels)
if err != nil {
t.Fatalf("json.Marshal(defaultModels) error = %v", err)
}
channelTemplateJSON, err := json.Marshal(providerManifest.ChannelTemplate)
if err != nil {
t.Fatalf("json.Marshal(channelTemplate) error = %v", err)
}
providerPK, err := store.Providers().Create(context.Background(), sqlite.Provider{
PackID: packPK,
ProviderID: providerManifest.ProviderID,
DisplayName: providerManifest.DisplayName,
BaseURL: providerManifest.BaseURL,
Platform: providerManifest.Platform,
AccountType: providerManifest.AccountType,
DefaultModelsJSON: string(defaultModelsJSON),
SmokeTestModel: providerManifest.SmokeTestModel,
ChannelTemplateJSON: string(channelTemplateJSON),
ManifestJSON: string(manifestJSON),
})
if err != nil {
t.Fatalf("Providers().Create() error = %v", err)
}
return packPK, providerPK
}
func mustCreateLegacyReusableBatch(t *testing.T, store *sqlite.DB, hostPK int64, packPK int64, providerPK int64, apiKey string, accountID string) int64 {
t.Helper()
batchID, err := store.ImportBatches().Create(context.Background(), sqlite.ImportBatch{
HostID: hostPK,
PackID: packPK,
ProviderID: providerPK,
Mode: "strict",
BatchStatus: "succeeded",
AccessStatus: "self_service_ready",
})
if err != nil {
t.Fatalf("ImportBatches().Create() error = %v", err)
}
if _, err := store.ImportBatchItems().Create(context.Background(), sqlite.ImportBatchItem{
BatchID: batchID,
KeyFingerprint: fullSHA256Fingerprint(apiKey),
AccountStatus: "passed",
ProbeSummaryJSON: fmt.Sprintf(`{"account_id":"%s"}`, accountID),
}); err != nil {
t.Fatalf("ImportBatchItems().Create() error = %v", err)
}
if _, err := store.ManagedResources().Create(context.Background(), sqlite.ManagedResource{
BatchID: batchID,
HostID: hostPK,
ResourceType: "account",
HostResourceID: accountID,
ResourceName: "legacy-account",
}); err != nil {
t.Fatalf("ManagedResources().Create(account) error = %v", err)
}
return batchID
}
func fullSHA256Fingerprint(value string) string {
sum := sha256.Sum256([]byte(strings.TrimSpace(value)))
return fmt.Sprintf("sha256:%x", sum[:])
}