327 lines
14 KiB
Go
327 lines
14 KiB
Go
package httpapi
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/identity"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/orchestration"
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/publicapi"
|
|
)
|
|
|
|
type httpJobStore struct {
|
|
values map[string]jobs.Job
|
|
claimed []jobs.Job
|
|
}
|
|
|
|
func (s *httpJobStore) ListJobs(_ context.Context, f jobs.ListFilter) ([]jobs.Job, error) {
|
|
out := []jobs.Job{}
|
|
for _, j := range s.values {
|
|
if f.Scope.Owns(j) && (f.Capability == "" || j.Capability == f.Capability) && (f.Status == "" || j.Status == f.Status) {
|
|
out = append(out, j)
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|
|
func (s *httpJobStore) FindJob(_ context.Context, sc jobs.Scope, id string) (jobs.Job, bool, error) {
|
|
j, ok := s.values[id]
|
|
return j, ok && sc.Owns(j), nil
|
|
}
|
|
func (s *httpJobStore) FindIdempotentJob(_ context.Context, sc jobs.Scope, key string) (jobs.Job, bool, error) {
|
|
for _, j := range s.values {
|
|
if sc.Owns(j) && j.IdempotencyKey == key {
|
|
return j, true, nil
|
|
}
|
|
}
|
|
return jobs.Job{}, false, nil
|
|
}
|
|
func (s *httpJobStore) CreateJob(_ context.Context, j jobs.Job) (jobs.Job, error) {
|
|
s.values[j.ID] = j
|
|
return j, nil
|
|
}
|
|
func (s *httpJobStore) UpdateJob(_ context.Context, id string, p jobs.Patch) (jobs.Job, error) {
|
|
j := s.values[id]
|
|
if p.Status != nil {
|
|
j.Status = *p.Status
|
|
}
|
|
if p.CompletedAt != nil {
|
|
j.CompletedAt = p.CompletedAt
|
|
}
|
|
if p.ClearLease {
|
|
j.LockedAt = nil
|
|
j.LockedBy = ""
|
|
}
|
|
s.values[id] = j
|
|
return j, nil
|
|
}
|
|
func (s *httpJobStore) DeleteJob(_ context.Context, id string) error {
|
|
delete(s.values, id)
|
|
return nil
|
|
}
|
|
func (s *httpJobStore) ClaimJobs(context.Context, string, int, int) ([]jobs.Job, error) {
|
|
return s.claimed, nil
|
|
}
|
|
|
|
type jobBuilderStub struct{ next int }
|
|
|
|
func (b *jobBuilderStub) Build(_ context.Context, input JobBuildInput) (jobs.CreateCommand, error) {
|
|
b.next++
|
|
raw, _ := json.Marshal(input.Body)
|
|
return jobs.CreateCommand{Job: jobs.Job{ID: "new-" + string(rune('0'+b.next)), OwnerID: input.Scope.OwnerID, ExternalClientID: input.Scope.ExternalClientID, Capability: input.Capability, Provider: "mock", ReqKey: "fixture", Status: jobs.StatusQueued, Prompt: "hello", RequestPayload: raw, IdempotencyKey: input.IdempotencyKey}, IdempotencyBody: input.Body}, nil
|
|
}
|
|
|
|
func newJobsHTTP(t *testing.T) (http.Handler, *httpJobStore) {
|
|
t.Helper()
|
|
store := &httpJobStore{values: map[string]jobs.Job{}}
|
|
platform, _ := NewPlatformAuthorizer(AuthState{}, nil)
|
|
service := jobs.NewService(store, func() time.Time { return time.Date(2026, 8, 13, 9, 0, 0, 0, time.UTC) })
|
|
builder := &jobBuilderStub{}
|
|
creation := &passthroughCreationCoordinator{service: service, builder: builder}
|
|
h, err := NewJobsHandler(JobsDependencies{Service: service, Platform: platform, Public: publicapi.NewAuthenticator(publicapi.Config{APIKeys: "agent-a:secret", InternalWorkerToken: "worker-secret", Production: true}), Builder: builder, Creation: creation}, JobsConfig{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return h, store
|
|
}
|
|
|
|
func TestJobsHTTPPlatformImageVideoLifecycle(t *testing.T) {
|
|
h, store := newJobsHTTP(t)
|
|
for _, path := range []string{"/api/generations/image", "/api/generations/video"} {
|
|
w := httptest.NewRecorder()
|
|
r := httptest.NewRequest(http.MethodPost, path, bytes.NewBufferString(`{"prompt":"hello"}`))
|
|
h.ServeHTTP(w, r)
|
|
if w.Code != 202 {
|
|
t.Fatalf("POST %s=%d %s", path, w.Code, w.Body.String())
|
|
}
|
|
}
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/api/generations/image", nil))
|
|
if w.Code != 200 {
|
|
t.Fatal(w.Code)
|
|
}
|
|
store.values["failed"] = jobs.Job{ID: "failed", OwnerID: "demo-merchant", Capability: "image.generate", Provider: "mock", ReqKey: "fixture", Status: jobs.StatusFailed, RequestPayload: []byte(`{}`)}
|
|
w = httptest.NewRecorder()
|
|
h.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/generations/image/failed/retry", nil))
|
|
if w.Code != 202 {
|
|
t.Fatalf("retry=%d %s", w.Code, w.Body.String())
|
|
}
|
|
w = httptest.NewRecorder()
|
|
h.ServeHTTP(w, httptest.NewRequest(http.MethodDelete, "/api/generations/image/failed", nil))
|
|
if w.Code != 200 {
|
|
t.Fatalf("delete=%d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestJobsHTTPPlatformRetryUsesFullCreationCoordinatorForActiveOwnedImage(t *testing.T) {
|
|
store := &httpJobStore{values: map[string]jobs.Job{"active": {ID: "active", OwnerID: "demo-merchant", Capability: "image.generate", Status: jobs.StatusRunning}}}
|
|
platform, _ := NewPlatformAuthorizer(AuthState{}, nil)
|
|
creation := &retryCreationCoordinatorStub{creationCoordinatorStub: creationCoordinatorStub{job: jobs.Job{ID: "fresh", RetryOf: "active", Status: jobs.StatusQueued}}}
|
|
h, err := NewJobsHandler(JobsDependencies{Service: jobs.NewService(store, time.Now), Platform: platform, Public: publicapi.NewAuthenticator(publicapi.Config{APIKeys: "agent-a:secret"}), Builder: &jobBuilderStub{}, Creation: creation}, JobsConfig{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/generations/image/active/retry", nil))
|
|
if w.Code != 202 || creation.retryCalls != 1 || creation.original.ID != "active" {
|
|
t.Fatalf("retry=%d %s calls=%d original=%#v", w.Code, w.Body.String(), creation.retryCalls, creation.original)
|
|
}
|
|
}
|
|
|
|
func TestJobsHTTPPublicScopeIdempotencyCancelAndMethods(t *testing.T) {
|
|
h, store := newJobsHTTP(t)
|
|
body := `{"capability":"image.generate","prompt":"hello"}`
|
|
request := func(method, path string) *httptest.ResponseRecorder {
|
|
r := httptest.NewRequest(method, path, bytes.NewBufferString(body))
|
|
r.Header.Set("Authorization", "Bearer secret")
|
|
r.Header.Set("Idempotency-Key", "same")
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, r)
|
|
return w
|
|
}
|
|
if w := request(http.MethodPost, "/api/v1/jobs"); w.Code != 202 {
|
|
t.Fatalf("create=%d %s", w.Code, w.Body.String())
|
|
}
|
|
if w := request(http.MethodPost, "/api/v1/jobs"); w.Code != 200 {
|
|
t.Fatalf("replay=%d %s", w.Code, w.Body.String())
|
|
}
|
|
store.values["other"] = jobs.Job{ID: "other", OwnerID: "api:agent-b", ExternalClientID: "agent-b", Status: jobs.StatusRunning}
|
|
if w := request(http.MethodGet, "/api/v1/jobs/other"); w.Code != 404 {
|
|
t.Fatalf("cross scope=%d", w.Code)
|
|
}
|
|
id := "new-1"
|
|
if w := request(http.MethodPost, "/api/v1/jobs/"+id+"/cancel"); w.Code != 200 {
|
|
t.Fatalf("cancel=%d %s", w.Code, w.Body.String())
|
|
}
|
|
if w := request(http.MethodDelete, "/api/v1/jobs"); w.Code != 405 || w.Header().Get("Allow") != "GET, POST" {
|
|
t.Fatalf("method=%d allow=%q", w.Code, w.Header().Get("Allow"))
|
|
}
|
|
}
|
|
|
|
func TestJobsHTTPPublicRejectsRelativeWebhookURL(t *testing.T) {
|
|
h, _ := newJobsHTTP(t)
|
|
for _, webhook := range []string{"/internal/callback", "ftp://example.test/hook", "https://user:pass@example.test/hook"} {
|
|
r := httptest.NewRequest(http.MethodPost, "/api/v1/jobs", bytes.NewBufferString(`{"capability":"image.generate","prompt":"hello","webhookUrl":"`+webhook+`"}`))
|
|
r.Header.Set("Authorization", "Bearer secret")
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, r)
|
|
if w.Code != http.StatusBadRequest {
|
|
t.Fatalf("webhook=%q status=%d body=%s", webhook, w.Code, w.Body.String())
|
|
}
|
|
}
|
|
r := httptest.NewRequest(http.MethodPost, "/api/v1/jobs", bytes.NewBufferString(`{"capability":"image.generate","prompt":"hello","webhookUrl":"https://hooks.example.test/done"}`))
|
|
r.Header.Set("Authorization", "Bearer secret")
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, r)
|
|
if w.Code != http.StatusAccepted {
|
|
t.Fatalf("valid status=%d body=%s", w.Code, w.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestJobsHTTPClampsPriorityBeforePassingTrustedBoundary(t *testing.T) {
|
|
store := &httpJobStore{values: map[string]jobs.Job{}}
|
|
platform, _ := NewPlatformAuthorizer(AuthState{}, nil)
|
|
creation := &creationCoordinatorStub{job: jobs.Job{ID: "created", Status: jobs.StatusQueued}}
|
|
h, _ := NewJobsHandler(JobsDependencies{Service: jobs.NewService(store, time.Now), Platform: platform, Public: publicapi.NewAuthenticator(publicapi.Config{APIKeys: "agent-a:secret"}), Builder: &jobBuilderStub{}, Creation: creation}, JobsConfig{})
|
|
r := httptest.NewRequest(http.MethodPost, "/api/v1/jobs", bytes.NewBufferString(`{"capability":"image.generate","prompt":"hello","priority":999}`))
|
|
r.Header.Set("Authorization", "Bearer secret")
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, r)
|
|
if w.Code != http.StatusAccepted || creation.publicInput.Body["priority"] != float64(100) {
|
|
t.Fatalf("status=%d input=%#v", w.Code, creation.publicInput.Body)
|
|
}
|
|
}
|
|
|
|
func TestJobsHTTPWorkerTickUsesWorkerAuth(t *testing.T) {
|
|
h, _ := newJobsHTTP(t)
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/internal/worker/tick", bytes.NewBufferString(`{}`)))
|
|
if w.Code != 401 {
|
|
t.Fatalf("unauthorized=%d", w.Code)
|
|
}
|
|
}
|
|
|
|
func TestJobsHTTPPlatformCreationPassesRefreshedSessionToCoordinator(t *testing.T) {
|
|
store := &httpJobStore{values: map[string]jobs.Job{}}
|
|
platform, _ := NewPlatformAuthorizer(AuthState{}, nil)
|
|
creation := &creationCoordinatorStub{job: jobs.Job{ID: "coordinated", Status: jobs.StatusQueued}}
|
|
h, err := NewJobsHandler(JobsDependencies{Service: jobs.NewService(store, time.Now), Platform: platform, Public: publicapi.NewAuthenticator(publicapi.Config{APIKeys: "agent-a:secret"}), Builder: &jobBuilderStub{}, Creation: creation}, JobsConfig{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/generations/image", bytes.NewBufferString(`{"prompt":"hello"}`)))
|
|
if w.Code != 202 || creation.platform.User.ID != "demo-merchant" || creation.platform.User.OrganizationID != "org-demo" {
|
|
t.Fatalf("response=%d %s session=%#v", w.Code, w.Body.String(), creation.platform)
|
|
}
|
|
}
|
|
|
|
func TestJobsHTTPMapsBillingChargeErrorsAndPublicUsesPublicCreation(t *testing.T) {
|
|
store := &httpJobStore{values: map[string]jobs.Job{}}
|
|
platform, _ := NewPlatformAuthorizer(AuthState{}, nil)
|
|
creation := &creationCoordinatorStub{err: &billing.StatusError{Status: 402, Err: billing.ErrInsufficientBalance}}
|
|
h, _ := NewJobsHandler(JobsDependencies{Service: jobs.NewService(store, time.Now), Platform: platform, Public: publicapi.NewAuthenticator(publicapi.Config{APIKeys: "agent-a:secret"}), Builder: &jobBuilderStub{}, Creation: creation}, JobsConfig{})
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/generations/image", bytes.NewBufferString(`{"prompt":"hello"}`)))
|
|
if w.Code != 402 || w.Body.String() != "{\"error\":\"生成任务计费失败。\"}\n" {
|
|
t.Fatalf("platform=%d %s", w.Code, w.Body.String())
|
|
}
|
|
creation.err = nil
|
|
r := httptest.NewRequest(http.MethodPost, "/api/v1/jobs", bytes.NewBufferString(`{"prompt":"hello"}`))
|
|
r.Header.Set("Authorization", "Bearer secret")
|
|
w = httptest.NewRecorder()
|
|
h.ServeHTTP(w, r)
|
|
if w.Code != 202 || creation.publicCalls != 1 || creation.platformCalls != 1 {
|
|
t.Fatalf("public=%d calls=%d/%d", w.Code, creation.platformCalls, creation.publicCalls)
|
|
}
|
|
}
|
|
|
|
func TestJobsHTTPMapsAllCreationBillingOutcomesWithoutLeakingErrors(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
err error
|
|
want int
|
|
}{
|
|
{name: "insufficient", err: &billing.StatusError{Status: 402, Err: billing.ErrInsufficientBalance}, want: 402},
|
|
{name: "idempotency", err: &billing.StatusError{Status: 409, Err: billing.ErrIdempotencyConflict}, want: 409},
|
|
{name: "unknown", err: errors.New("postgres password leaked"), want: 500},
|
|
}
|
|
for _, test := range tests {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
store := &httpJobStore{values: map[string]jobs.Job{}}
|
|
platform, _ := NewPlatformAuthorizer(AuthState{}, nil)
|
|
creation := &creationCoordinatorStub{err: test.err}
|
|
h, _ := NewJobsHandler(JobsDependencies{Service: jobs.NewService(store, time.Now), Platform: platform, Public: publicapi.NewAuthenticator(publicapi.Config{APIKeys: "agent-a:secret"}), Builder: &jobBuilderStub{}, Creation: creation}, JobsConfig{})
|
|
w := httptest.NewRecorder()
|
|
h.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/generations/image", bytes.NewBufferString(`{"prompt":"hello"}`)))
|
|
if w.Code != test.want || strings.Contains(w.Body.String(), "postgres") {
|
|
t.Fatalf("status=%d body=%s", w.Code, w.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
type creationCoordinatorStub struct {
|
|
job jobs.Job
|
|
err error
|
|
platform identity.Session
|
|
platformCalls, publicCalls int
|
|
publicInput orchestration.CreationInput
|
|
}
|
|
|
|
type retryCreationCoordinatorStub struct {
|
|
creationCoordinatorStub
|
|
retryCalls int
|
|
original jobs.Job
|
|
}
|
|
|
|
func (s *retryCreationCoordinatorStub) RetryPlatform(_ context.Context, _ identity.Session, original jobs.Job) (jobs.Job, error) {
|
|
s.retryCalls++
|
|
s.original = original
|
|
return s.job, s.err
|
|
}
|
|
|
|
type passthroughCreationCoordinator struct {
|
|
service *jobs.Service
|
|
builder JobBuilder
|
|
}
|
|
|
|
func (s *passthroughCreationCoordinator) create(ctx context.Context, scope jobs.Scope, in orchestration.CreationInput) (jobs.Job, bool, error) {
|
|
command, err := s.builder.Build(ctx, JobBuildInput{Scope: scope, Capability: in.Capability, Body: in.Body, IdempotencyKey: in.IdempotencyKey})
|
|
if err != nil {
|
|
return jobs.Job{}, false, err
|
|
}
|
|
return s.service.Create(ctx, command)
|
|
}
|
|
func (s *passthroughCreationCoordinator) CreatePlatform(ctx context.Context, session identity.Session, in orchestration.CreationInput) (jobs.Job, bool, error) {
|
|
return s.create(ctx, jobs.Scope{OwnerID: session.User.ID}, in)
|
|
}
|
|
func (s *passthroughCreationCoordinator) CreatePublic(ctx context.Context, in orchestration.CreationInput) (jobs.Job, bool, error) {
|
|
return s.create(ctx, jobs.Scope{OwnerID: in.OwnerID, ExternalClientID: in.ExternalClientID}, in)
|
|
}
|
|
func (s *passthroughCreationCoordinator) RetryPlatform(ctx context.Context, session identity.Session, original jobs.Job) (jobs.Job, error) {
|
|
body := map[string]any{"prompt": original.Prompt}
|
|
job, _, err := s.create(ctx, jobs.Scope{OwnerID: session.User.ID}, orchestration.CreationInput{Capability: original.Capability, Body: body})
|
|
job.RetryOf = original.ID
|
|
return job, err
|
|
}
|
|
|
|
func (s *creationCoordinatorStub) CreatePlatform(_ context.Context, session identity.Session, _ orchestration.CreationInput) (jobs.Job, bool, error) {
|
|
s.platform = session
|
|
s.platformCalls++
|
|
return s.job, false, s.err
|
|
}
|
|
func (s *creationCoordinatorStub) CreatePublic(_ context.Context, input orchestration.CreationInput) (jobs.Job, bool, error) {
|
|
s.publicCalls++
|
|
s.publicInput = input
|
|
return s.job, false, s.err
|
|
}
|