700 lines
25 KiB
Go
700 lines
25 KiB
Go
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"
|
|
)
|
|
|
|
func TestBatchImportHTTP(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("POST create run returns run summary", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
handler := NewAPIHandler("secret-token", ActionSet{
|
|
CreateBatchImportRun: func(_ context.Context, req CreateBatchImportRunRequest) (BatchImportRunCreateResponse, error) {
|
|
if req.HostID != "host-1" {
|
|
t.Fatalf("HostID = %q, want host-1", req.HostID)
|
|
}
|
|
if req.AccessMode != "subscription" {
|
|
t.Fatalf("AccessMode = %q, want subscription", req.AccessMode)
|
|
}
|
|
if len(req.SubscriptionUsers) != 1 || req.SubscriptionUsers[0] != "user-1" {
|
|
t.Fatalf("SubscriptionUsers = %#v, want [user-1]", req.SubscriptionUsers)
|
|
}
|
|
if len(req.Entries) != 1 || req.Entries[0].BaseURL != "https://kimi.example.com/v1" {
|
|
t.Fatalf("Entries = %#v, want request payload", req.Entries)
|
|
}
|
|
return BatchImportRunCreateResponse{
|
|
RunID: "run_20260522_0001",
|
|
State: "running",
|
|
ResultPage: "/batch-import/runs/run_20260522_0001",
|
|
TotalItems: 1,
|
|
ActiveItems: 0,
|
|
DegradedItems: 0,
|
|
BrokenItems: 0,
|
|
WarningItems: 0,
|
|
}, nil
|
|
},
|
|
})
|
|
|
|
req := httptestRequest(t, http.MethodPost, "/api/batch-import/runs", map[string]any{
|
|
"host_id": "host-1",
|
|
"mode": "strict",
|
|
"access_mode": "subscription",
|
|
"subscription_users": []string{"user-1"},
|
|
"subscription_days": 30,
|
|
"entries": []map[string]any{
|
|
{"base_url": "https://kimi.example.com/v1", "api_key": "sk-test", "requested_models": []string{"kimi-k2.6"}},
|
|
},
|
|
}, "secret-token")
|
|
res := httptestRecorder(handler, req)
|
|
assertStatusCode(t, res, http.StatusOK)
|
|
assertJSONContains(t, res.Body().Bytes(), "run_id", "run_20260522_0001")
|
|
assertJSONContains(t, res.Body().Bytes(), "result_page", "/batch-import/runs/run_20260522_0001")
|
|
})
|
|
|
|
t.Run("subscription request requires subscription fields", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
handler := NewAPIHandler("secret-token", ActionSet{
|
|
CreateBatchImportRun: func(_ context.Context, req CreateBatchImportRunRequest) (BatchImportRunCreateResponse, error) {
|
|
t.Fatal("CreateBatchImportRun should not be called when request is invalid")
|
|
return BatchImportRunCreateResponse{}, nil
|
|
},
|
|
})
|
|
|
|
req := httptestRequest(t, http.MethodPost, "/api/batch-import/runs", map[string]any{
|
|
"host_id": "host-1",
|
|
"mode": "strict",
|
|
"access_mode": "subscription",
|
|
"entries": []map[string]any{
|
|
{"base_url": "https://kimi.example.com/v1", "api_key": "sk-test"},
|
|
},
|
|
}, "secret-token")
|
|
res := httptestRecorder(handler, req)
|
|
assertStatusCode(t, res, http.StatusBadRequest)
|
|
assertJSONContains(t, res.Body().Bytes(), "error.code", "invalid_request")
|
|
})
|
|
|
|
t.Run("self service request requires probe api key", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
handler := NewAPIHandler("secret-token", ActionSet{
|
|
CreateBatchImportRun: func(_ context.Context, req CreateBatchImportRunRequest) (BatchImportRunCreateResponse, error) {
|
|
t.Fatal("CreateBatchImportRun should not be called when request is invalid")
|
|
return BatchImportRunCreateResponse{}, nil
|
|
},
|
|
})
|
|
|
|
req := httptestRequest(t, http.MethodPost, "/api/batch-import/runs", map[string]any{
|
|
"host_id": "host-1",
|
|
"mode": "partial",
|
|
"access_mode": "self_service",
|
|
"entries": []map[string]any{
|
|
{"base_url": "https://deepseek.example.com/v1", "api_key": "sk-test"},
|
|
},
|
|
}, "secret-token")
|
|
res := httptestRecorder(handler, req)
|
|
assertStatusCode(t, res, http.StatusBadRequest)
|
|
assertJSONContains(t, res.Body().Bytes(), "error.code", "invalid_request")
|
|
assertJSONContains(t, res.Body().Bytes(), "error.message", "probe_api_key is required when access_mode=self_service")
|
|
})
|
|
|
|
t.Run("create run requires host id", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
handler := NewAPIHandler("secret-token", ActionSet{
|
|
CreateBatchImportRun: func(_ context.Context, req CreateBatchImportRunRequest) (BatchImportRunCreateResponse, error) {
|
|
t.Fatal("CreateBatchImportRun should not be called when host_id is missing")
|
|
return BatchImportRunCreateResponse{}, nil
|
|
},
|
|
})
|
|
|
|
req := httptestRequest(t, http.MethodPost, "/api/batch-import/runs", map[string]any{
|
|
"mode": "strict",
|
|
"access_mode": "subscription",
|
|
"subscription_users": []string{"user-1"},
|
|
"subscription_days": 30,
|
|
"entries": []map[string]any{
|
|
{"base_url": "https://kimi.example.com/v1", "api_key": "sk-test"},
|
|
},
|
|
}, "secret-token")
|
|
res := httptestRecorder(handler, req)
|
|
assertStatusCode(t, res, http.StatusBadRequest)
|
|
assertJSONContains(t, res.Body().Bytes(), "error.code", "invalid_request")
|
|
assertJSONContains(t, res.Body().Bytes(), "error.message", "host_id is required")
|
|
})
|
|
|
|
t.Run("create run action wires real batch pipeline", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
server := httptest.NewServer(newBatchImportActionStubServer(t))
|
|
defer server.Close()
|
|
|
|
dsn := testutil.SQLiteTestDSN(t, "state.db", true)
|
|
store := testutil.OpenSQLiteStore(t, dsn)
|
|
defer closeAppTestStore(t, store)
|
|
|
|
if _, err := store.Hosts().Create(context.Background(), sqlite.Host{
|
|
HostID: "host-1",
|
|
BaseURL: server.URL,
|
|
HostVersion: "0.1.126",
|
|
CapabilityProbeJSON: "{}",
|
|
AuthType: "apikey",
|
|
AuthToken: "host-token",
|
|
}); err != nil {
|
|
t.Fatalf("Hosts().Create() error = %v", err)
|
|
}
|
|
|
|
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() error = %v", err)
|
|
}
|
|
if result.State != string("completed") {
|
|
t.Fatalf("result.State = %q, want completed", result.State)
|
|
}
|
|
if result.ActiveItems != 1 {
|
|
t.Fatalf("result.ActiveItems = %d, want 1", result.ActiveItems)
|
|
}
|
|
|
|
run, err := store.ImportRuns().GetByRunID(context.Background(), result.RunID)
|
|
if err != nil {
|
|
t.Fatalf("ImportRuns().GetByRunID() error = %v", err)
|
|
}
|
|
if run.State != "completed" {
|
|
t.Fatalf("run.State = %q, want completed", run.State)
|
|
}
|
|
items, err := store.ImportRunItems().ListByRunID(context.Background(), result.RunID)
|
|
if err != nil {
|
|
t.Fatalf("ImportRunItems().ListByRunID() error = %v", err)
|
|
}
|
|
if len(items) != 1 {
|
|
t.Fatalf("len(items) = %d, want 1", len(items))
|
|
}
|
|
if items[0].CurrentStage != "done" || items[0].AccessStatus != "active" {
|
|
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) {
|
|
t.Parallel()
|
|
|
|
t.Run("handleCreateBatchImportRun requires action", func(t *testing.T) {
|
|
t.Parallel()
|
|
req := httptestRequest(t, http.MethodPost, "/api/batch-import/runs", map[string]any{}, "")
|
|
rec := &responseRecorder{header: map[string][]string{}}
|
|
handleCreateBatchImportRun(rec, req, nil)
|
|
assertStatusCode(t, rec, http.StatusInternalServerError)
|
|
assertJSONContains(t, rec.Body().Bytes(), "error.code", "server_misconfigured")
|
|
})
|
|
|
|
t.Run("handleCreateBatchImportRun classifies action error", func(t *testing.T) {
|
|
t.Parallel()
|
|
req := httptestRequest(t, http.MethodPost, "/api/batch-import/runs", map[string]any{
|
|
"host_id": "host-1",
|
|
"mode": "strict",
|
|
"access_mode": "self_service",
|
|
"probe_api_key": "probe-key",
|
|
"entries": []map[string]any{
|
|
{"base_url": "https://kimi.example.com/v1", "api_key": "sk-test"},
|
|
},
|
|
}, "")
|
|
rec := &responseRecorder{header: map[string][]string{}}
|
|
handleCreateBatchImportRun(rec, req, func(context.Context, CreateBatchImportRunRequest) (BatchImportRunCreateResponse, error) {
|
|
return BatchImportRunCreateResponse{}, fmt.Errorf("host x not found")
|
|
})
|
|
assertStatusCode(t, rec, http.StatusNotFound)
|
|
assertJSONContains(t, rec.Body().Bytes(), "error.code", "not_found")
|
|
})
|
|
|
|
t.Run("handleListBatchImportRuns requires action", func(t *testing.T) {
|
|
t.Parallel()
|
|
req := httptestRequest(t, http.MethodGet, "/api/batch-import/runs", nil, "")
|
|
rec := &responseRecorder{header: map[string][]string{}}
|
|
handleListBatchImportRuns(rec, req, nil)
|
|
assertStatusCode(t, rec, http.StatusInternalServerError)
|
|
assertJSONContains(t, rec.Body().Bytes(), "error.code", "server_misconfigured")
|
|
})
|
|
|
|
t.Run("handleListBatchImportRuns returns empty array", func(t *testing.T) {
|
|
t.Parallel()
|
|
req := httptestRequest(t, http.MethodGet, "/api/batch-import/runs?limit=5", nil, "")
|
|
rec := &responseRecorder{header: map[string][]string{}}
|
|
handleListBatchImportRuns(rec, req, func(_ context.Context, got ListBatchImportRunsRequest) (ListBatchImportRunsResponse, error) {
|
|
if got.Limit != 5 {
|
|
t.Fatalf("ListBatchImportRunsRequest.Limit = %d, want 5", got.Limit)
|
|
}
|
|
return ListBatchImportRunsResponse{}, nil
|
|
})
|
|
assertStatusCode(t, rec, http.StatusOK)
|
|
runs, ok := decodeTopLevelArray(t, rec.Body().Bytes(), "runs")
|
|
if !ok || len(runs) != 0 {
|
|
t.Fatalf("runs = %#v, want empty array", runs)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestBatchImportRejectsOversizedJSONBody(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
handler := NewAPIHandler("secret-token", ActionSet{
|
|
CreateBatchImportRun: func(_ context.Context, req CreateBatchImportRunRequest) (BatchImportRunCreateResponse, error) {
|
|
t.Fatal("CreateBatchImportRun should not be called for oversized body")
|
|
return BatchImportRunCreateResponse{}, nil
|
|
},
|
|
})
|
|
|
|
payload := `{"host_id":"host-1","mode":"strict","access_mode":"self_service","probe_api_key":"probe-key","entries":[{"base_url":"https://kimi.example.com/v1","api_key":"` + strings.Repeat("x", int(maxJSONBodyBytes)) + `"}]}`
|
|
req := httptest.NewRequest(http.MethodPost, "/api/batch-import/runs", strings.NewReader(payload))
|
|
req.Header.Set("Authorization", "Bearer secret-token")
|
|
res := httptestRecorder(handler, req)
|
|
assertStatusCode(t, res, http.StatusRequestEntityTooLarge)
|
|
assertJSONContains(t, res.Body().Bytes(), "error.code", "request_too_large")
|
|
}
|
|
|
|
func newBatchImportActionStubServer(t *testing.T) http.Handler {
|
|
t.Helper()
|
|
|
|
mux := http.NewServeMux()
|
|
mux.HandleFunc("/api/v1/admin/system/version", func(w http.ResponseWriter, r *http.Request) {
|
|
if !requireBatchImportActionAdminToken(t, w, r) {
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, map[string]any{"data": map[string]any{"version": "0.1.126"}})
|
|
})
|
|
mux.HandleFunc("/api/v1/admin/groups", func(w http.ResponseWriter, r *http.Request) {
|
|
if !requireBatchImportActionAdminToken(t, w, r) {
|
|
return
|
|
}
|
|
switch r.Method {
|
|
case http.MethodGet:
|
|
writeJSON(w, http.StatusOK, map[string]any{"data": []map[string]any{}})
|
|
case http.MethodPost:
|
|
writeJSON(w, http.StatusCreated, map[string]any{"data": map[string]any{"id": "group_1", "name": "batch-import-group"}})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
})
|
|
mux.HandleFunc("/api/v1/admin/channels", func(w http.ResponseWriter, r *http.Request) {
|
|
if !requireBatchImportActionAdminToken(t, w, r) {
|
|
return
|
|
}
|
|
switch r.Method {
|
|
case http.MethodGet:
|
|
writeJSON(w, http.StatusOK, map[string]any{"data": []map[string]any{}})
|
|
case http.MethodPost:
|
|
writeJSON(w, http.StatusCreated, map[string]any{"data": map[string]any{"id": "channel_1", "name": "batch-import-channel"}})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
})
|
|
mux.HandleFunc("/api/v1/admin/payment/plans", func(w http.ResponseWriter, r *http.Request) {
|
|
if !requireBatchImportActionAdminToken(t, w, r) {
|
|
return
|
|
}
|
|
switch r.Method {
|
|
case http.MethodGet:
|
|
writeJSON(w, http.StatusOK, map[string]any{"data": []map[string]any{}})
|
|
case http.MethodPost:
|
|
writeJSON(w, http.StatusCreated, map[string]any{"data": map[string]any{"id": "plan_1", "name": "batch-import-plan"}})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
})
|
|
mux.HandleFunc("/api/v1/admin/accounts", 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{}, "pages": 1}})
|
|
})
|
|
mux.HandleFunc("/api/v1/admin/accounts/batch", func(w http.ResponseWriter, r *http.Request) {
|
|
if !requireBatchImportActionAdminToken(t, w, r) {
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, map[string]any{"data": []map[string]any{{"id": "account_1", "name": "batch-import-01"}}})
|
|
})
|
|
mux.HandleFunc("/api/v1/admin/accounts/__probe__/test", func(w http.ResponseWriter, r *http.Request) {
|
|
if !requireBatchImportActionAdminToken(t, w, r) {
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "probe only"})
|
|
})
|
|
mux.HandleFunc("/api/v1/admin/accounts/__probe__/models", func(w http.ResponseWriter, r *http.Request) {
|
|
if !requireBatchImportActionAdminToken(t, w, r) {
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "probe only"})
|
|
})
|
|
mux.HandleFunc("/api/v1/admin/accounts/account_1/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/account_1/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/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
|
|
}
|
|
switch r.Method {
|
|
case http.MethodGet:
|
|
writeJSON(w, http.StatusBadRequest, map[string]any{"error": "probe only"})
|
|
case http.MethodPost:
|
|
writeJSON(w, http.StatusOK, map[string]any{"data": map[string]any{"id": "subscription_1"}})
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
})
|
|
mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
|
|
switch strings.TrimSpace(r.Header.Get("Authorization")) {
|
|
case "Bearer entry-key", "Bearer gateway-key":
|
|
writeJSON(w, http.StatusOK, map[string]any{"data": []map[string]any{{"id": "kimi-k2.6"}}})
|
|
default:
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
|
|
}
|
|
})
|
|
mux.HandleFunc("/v1/responses", func(w http.ResponseWriter, r *http.Request) {
|
|
if strings.TrimSpace(r.Header.Get("Authorization")) != "Bearer entry-key" {
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
|
|
return
|
|
}
|
|
w.WriteHeader(http.StatusForbidden)
|
|
_, _ = w.Write([]byte(`{"error":"responses unsupported"}`))
|
|
})
|
|
mux.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
|
|
switch strings.TrimSpace(r.Header.Get("Authorization")) {
|
|
case "Bearer entry-key", "Bearer gateway-key":
|
|
writeJSON(w, http.StatusOK, map[string]any{
|
|
"id": "chatcmpl_batch_import",
|
|
"choices": []map[string]any{{
|
|
"index": 0,
|
|
"message": map[string]any{
|
|
"role": "assistant",
|
|
"content": "pong",
|
|
},
|
|
}},
|
|
})
|
|
default:
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
|
|
}
|
|
})
|
|
|
|
return mux
|
|
}
|
|
|
|
func requireBatchImportActionAdminToken(t *testing.T, w http.ResponseWriter, r *http.Request) bool {
|
|
t.Helper()
|
|
if strings.TrimSpace(r.Header.Get("x-api-key")) != "host-token" {
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
_, _ = w.Write([]byte(`{"error":"unauthorized"}`))
|
|
return false
|
|
}
|
|
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[:])
|
|
}
|