Files
sub2api/backend/internal/service/image_task_test.go
T
李建琦 6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
2026-08-21 18:30:13 +08:00

92 lines
2.9 KiB
Go

package service
import (
"context"
"encoding/json"
"errors"
"net/http"
"testing"
"time"
"github.com/stretchr/testify/require"
)
type imageTaskMemoryStore struct {
task *ImageTaskRecord
ttl time.Duration
saveErr error
getErr error
}
func (s *imageTaskMemoryStore) Save(_ context.Context, task *ImageTaskRecord, ttl time.Duration) error {
if s.saveErr != nil {
return s.saveErr
}
copy := *task
s.task = &copy
s.ttl = ttl
return nil
}
func (s *imageTaskMemoryStore) Get(_ context.Context, _ string) (*ImageTaskRecord, error) {
if s.getErr != nil {
return nil, s.getErr
}
if s.task == nil {
return nil, ErrImageTaskNotFound
}
copy := *s.task
return &copy, nil
}
func TestImageTaskServiceLifecycleAndOwnership(t *testing.T) {
store := &imageTaskMemoryStore{}
svc := NewImageTaskServiceWithOptions(store, time.Hour, 10*time.Minute)
owner := ImageTaskOwner{UserID: 7, APIKeyID: 9}
created, err := svc.Create(context.Background(), owner)
require.NoError(t, err)
require.Equal(t, ImageTaskStatusProcessing, created.Status)
require.Equal(t, created.ID, created.TaskID)
require.Equal(t, "image.generation.task", created.Object)
require.Equal(t, time.Hour, store.ttl)
require.Equal(t, owner.UserID, store.task.UserID)
require.Equal(t, owner.APIKeyID, store.task.APIKeyID)
_, err = svc.Get(context.Background(), ImageTaskOwner{UserID: 7, APIKeyID: 10}, created.ID)
require.ErrorIs(t, err, ErrImageTaskNotFound)
result := json.RawMessage(`{"created":123,"data":[{"url":"https://example.test/image.png"}]}`)
require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, result))
completed, err := svc.Get(context.Background(), owner, created.ID)
require.NoError(t, err)
require.Equal(t, ImageTaskStatusCompleted, completed.Status)
require.Equal(t, http.StatusOK, completed.HTTPStatus)
require.Equal(t, "https://example.test/image.png", completed.ImageURL)
require.JSONEq(t, string(result), string(completed.Result))
require.NotNil(t, completed.CompletedAt)
}
func TestImageTaskServiceInvalidResultBecomesFailed(t *testing.T) {
store := &imageTaskMemoryStore{}
svc := NewImageTaskServiceWithOptions(store, time.Hour, time.Minute)
created, err := svc.Create(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2})
require.NoError(t, err)
require.NoError(t, svc.Complete(context.Background(), created.ID, http.StatusOK, json.RawMessage(`not-json`)))
got, err := svc.Get(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2}, created.ID)
require.NoError(t, err)
require.Equal(t, ImageTaskStatusFailed, got.Status)
require.Equal(t, http.StatusBadGateway, got.HTTPStatus)
require.Contains(t, string(got.Error), "non-JSON")
}
func TestImageTaskServiceMapsStoreFailures(t *testing.T) {
store := &imageTaskMemoryStore{saveErr: errors.New("redis down")}
svc := NewImageTaskService(store)
_, err := svc.Create(context.Background(), ImageTaskOwner{UserID: 1, APIKeyID: 2})
require.ErrorIs(t, err, ErrImageTaskUnavailable)
}