Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）
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
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
This commit is contained in:
李建琦
@@ -0,0 +1,381 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha1"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newBatchImageRepositoryWithSQL(sqlq batchImageSQLExecutor) *batchImageRepository {
|
||||
return &batchImageRepository{sql: sqlq}
|
||||
}
|
||||
|
||||
func TestBatchImageRepository_CreateJobAndDuplicates(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := newBatchImageRepositoryWithSQL(tx)
|
||||
batchID := batchImageTestID(t, "create")
|
||||
|
||||
job, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{
|
||||
BatchID: batchID,
|
||||
UserID: 1001,
|
||||
Provider: service.BatchImageProviderGeminiAPI,
|
||||
Model: "gemini-2.5-flash-image",
|
||||
ItemCount: 2,
|
||||
EstimatedCost: 0.02,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, batchID, job.BatchID)
|
||||
require.Equal(t, service.BatchImageJobStatusCreated, job.Status)
|
||||
require.Equal(t, "USD", job.Currency)
|
||||
|
||||
_, err = repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{
|
||||
BatchID: batchID,
|
||||
UserID: 1001,
|
||||
Provider: service.BatchImageProviderGeminiAPI,
|
||||
Model: "gemini-2.5-flash-image",
|
||||
ItemCount: 1,
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.True(t, errors.Is(err, service.ErrBatchImageJobExists))
|
||||
}
|
||||
|
||||
func TestBatchImageRepository_InvalidProvider(t *testing.T) {
|
||||
tx := testTx(t)
|
||||
repo := newBatchImageRepositoryWithSQL(tx)
|
||||
|
||||
_, err := repo.CreateBatchImageJob(context.Background(), service.CreateBatchImageJobParams{
|
||||
BatchID: batchImageTestID(t, "provider"),
|
||||
UserID: 1001,
|
||||
Provider: "unknown",
|
||||
Model: "gemini-2.5-flash-image",
|
||||
ItemCount: 1,
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.True(t, errors.Is(err, service.ErrBatchImageInvalidProvider))
|
||||
}
|
||||
|
||||
func TestBatchImageRepository_TransitionIncrementsVersionAndEvents(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := newBatchImageRepositoryWithSQL(tx)
|
||||
batchID := batchImageTestID(t, "transition")
|
||||
now := time.Date(2026, 7, 3, 8, 0, 0, 0, time.UTC)
|
||||
|
||||
_, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{
|
||||
BatchID: batchID,
|
||||
UserID: 1001,
|
||||
Provider: service.BatchImageProviderVertex,
|
||||
Model: "gemini-2.5-flash-image",
|
||||
ItemCount: 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusUploading, service.BatchImageTransitionOptions{
|
||||
EventType: "status_changed",
|
||||
EventPayload: map[string]any{"to": service.BatchImageJobStatusUploading},
|
||||
Now: &now,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
job, err := repo.GetBatchImageJobByBatchID(ctx, batchID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, service.BatchImageJobStatusUploading, job.Status)
|
||||
require.Equal(t, 1, job.Version)
|
||||
|
||||
var eventCount int
|
||||
err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM batch_image_events WHERE job_id = $1 AND event_type = 'status_changed'`, batchID).Scan(&eventCount)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, eventCount)
|
||||
}
|
||||
|
||||
func TestBatchImageRepository_InvalidTransition(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := newBatchImageRepositoryWithSQL(tx)
|
||||
batchID := batchImageTestID(t, "invalid-transition")
|
||||
|
||||
_, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{
|
||||
BatchID: batchID,
|
||||
UserID: 1001,
|
||||
Provider: service.BatchImageProviderGeminiAPI,
|
||||
Model: "gemini-2.5-flash-image",
|
||||
ItemCount: 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusRunning, service.BatchImageTransitionOptions{})
|
||||
require.Error(t, err)
|
||||
require.True(t, errors.Is(err, service.ErrBatchImageInvalidTransition))
|
||||
}
|
||||
|
||||
func TestBatchImageRepository_TerminalStatusCannotMoveBack(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := newBatchImageRepositoryWithSQL(tx)
|
||||
batchID := batchImageTestID(t, "terminal")
|
||||
|
||||
_, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{
|
||||
BatchID: batchID,
|
||||
UserID: 1001,
|
||||
Provider: service.BatchImageProviderGeminiAPI,
|
||||
Model: "gemini-2.5-flash-image",
|
||||
Status: service.BatchImageJobStatusCompleted,
|
||||
ItemCount: 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusRunning, service.BatchImageTransitionOptions{})
|
||||
require.Error(t, err)
|
||||
require.True(t, errors.Is(err, service.ErrBatchImageInvalidTransition))
|
||||
}
|
||||
|
||||
func TestBatchImageRepository_ItemCustomIDUniqueness(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := newBatchImageRepositoryWithSQL(tx)
|
||||
firstBatchID := batchImageTestID(t, "items-a")
|
||||
secondBatchID := batchImageTestID(t, "items-b")
|
||||
|
||||
for _, batchID := range []string{firstBatchID, secondBatchID} {
|
||||
_, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{
|
||||
BatchID: batchID,
|
||||
UserID: 1001,
|
||||
Provider: service.BatchImageProviderGeminiAPI,
|
||||
Model: "gemini-2.5-flash-image",
|
||||
ItemCount: 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
_, err := repo.CreateBatchImageItem(ctx, service.CreateBatchImageItemParams{
|
||||
JobID: firstBatchID,
|
||||
CustomID: "line-1",
|
||||
Status: service.BatchImageItemStatusSuccess,
|
||||
ImageCount: 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
_, err = tx.ExecContext(ctx, `SAVEPOINT batch_image_duplicate_item`)
|
||||
require.NoError(t, err)
|
||||
_, err = repo.CreateBatchImageItem(ctx, service.CreateBatchImageItemParams{
|
||||
JobID: firstBatchID,
|
||||
CustomID: "line-1",
|
||||
Status: service.BatchImageItemStatusFailed,
|
||||
})
|
||||
require.Error(t, err)
|
||||
require.True(t, errors.Is(err, service.ErrBatchImageItemExists))
|
||||
_, rollbackErr := tx.ExecContext(ctx, `ROLLBACK TO SAVEPOINT batch_image_duplicate_item`)
|
||||
require.NoError(t, rollbackErr)
|
||||
|
||||
_, err = repo.CreateBatchImageItem(ctx, service.CreateBatchImageItemParams{
|
||||
JobID: secondBatchID,
|
||||
CustomID: "line-1",
|
||||
Status: service.BatchImageItemStatusSuccess,
|
||||
ImageCount: 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
items, err := repo.ListBatchImageItems(ctx, firstBatchID, service.BatchImageItemFilter{})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, items, 1)
|
||||
}
|
||||
|
||||
func TestBatchImageRepository_ReplaceBatchImageItemsForJob(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := newBatchImageRepositoryWithSQL(tx)
|
||||
batchID := batchImageTestID(t, "replace-items")
|
||||
lineOne := 1
|
||||
lineTwo := 2
|
||||
|
||||
_, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{
|
||||
BatchID: batchID,
|
||||
UserID: 1001,
|
||||
Provider: service.BatchImageProviderGeminiAPI,
|
||||
Model: "gemini-2.5-flash-image",
|
||||
ItemCount: 2,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// 非 indexing 状态不允许重建 item 表：防止锁过期后掉队的 worker
|
||||
// 重写已完成/已结算 job 的条目。
|
||||
err = repo.ReplaceBatchImageItemsForJob(ctx, batchID, []service.CreateBatchImageItemParams{
|
||||
{CustomID: "old", Status: service.BatchImageItemStatusSuccess, SourceLineNumber: &lineOne, ImageCount: 1},
|
||||
}, service.BatchImageCounts{SuccessCount: 1})
|
||||
require.ErrorIs(t, err, service.ErrBatchImageIndexStateConflict)
|
||||
|
||||
require.NoError(t, repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusSubmitted, service.BatchImageTransitionOptions{}))
|
||||
require.NoError(t, repo.TransitionBatchImageJobStatus(ctx, batchID, service.BatchImageJobStatusIndexing, service.BatchImageTransitionOptions{}))
|
||||
|
||||
err = repo.ReplaceBatchImageItemsForJob(ctx, batchID, []service.CreateBatchImageItemParams{
|
||||
{CustomID: "old", Status: service.BatchImageItemStatusSuccess, SourceLineNumber: &lineOne, ImageCount: 1},
|
||||
}, service.BatchImageCounts{SuccessCount: 1})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = repo.ReplaceBatchImageItemsForJob(ctx, batchID, []service.CreateBatchImageItemParams{
|
||||
{CustomID: "new-ok", Status: service.BatchImageItemStatusSuccess, SourceLineNumber: &lineOne, ImageCount: 1},
|
||||
{CustomID: "new-fail", Status: service.BatchImageItemStatusFailed, SourceLineNumber: &lineTwo, ErrorCode: batchImageTestStringPtr("SAFETY_BLOCKED")},
|
||||
}, service.BatchImageCounts{SuccessCount: 1, FailCount: 1})
|
||||
require.NoError(t, err)
|
||||
|
||||
items, err := repo.ListBatchImageItems(ctx, batchID, service.BatchImageItemFilter{})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, items, 2)
|
||||
require.Equal(t, "new-ok", items[0].CustomID)
|
||||
require.Equal(t, "new-fail", items[1].CustomID)
|
||||
|
||||
job, err := repo.GetBatchImageJobByBatchID(ctx, batchID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, job.SuccessCount)
|
||||
require.Equal(t, 1, job.FailCount)
|
||||
}
|
||||
|
||||
func TestBatchImageRepository_MarkBatchImageJobSettled(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := newBatchImageRepositoryWithSQL(tx)
|
||||
batchID := batchImageTestID(t, "settled")
|
||||
apiKeyID := int64(2001)
|
||||
accountID := int64(3001)
|
||||
providerJob := "providers/job"
|
||||
outputRef := "files/output"
|
||||
now := time.Date(2026, 7, 4, 10, 0, 0, 0, time.UTC)
|
||||
|
||||
_, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{
|
||||
BatchID: batchID,
|
||||
UserID: 1001,
|
||||
APIKeyID: &apiKeyID,
|
||||
AccountID: &accountID,
|
||||
Provider: service.BatchImageProviderGeminiAPI,
|
||||
Model: "gemini-image",
|
||||
Status: service.BatchImageJobStatusSettling,
|
||||
ProviderJobName: &providerJob,
|
||||
ProviderOutputRef: &outputRef,
|
||||
ItemCount: 3,
|
||||
SuccessCount: 2,
|
||||
FailCount: 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = repo.MarkBatchImageJobSettled(ctx, service.MarkBatchImageJobSettledParams{
|
||||
BatchID: batchID,
|
||||
ActualCost: 0.5,
|
||||
ManifestHash: "manifest-hash",
|
||||
EventPayload: map[string]any{"request_id": "batch_image_settlement:" + batchID},
|
||||
Now: &now,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
job, err := repo.GetBatchImageJobByBatchID(ctx, batchID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, service.BatchImageJobStatusCompleted, job.Status)
|
||||
require.NotNil(t, job.ActualCost)
|
||||
require.Equal(t, 0.5, *job.ActualCost)
|
||||
require.Equal(t, "manifest-hash", batchImageDerefTest(job.ManifestHash))
|
||||
require.NotNil(t, job.SettledAt)
|
||||
require.Equal(t, now, *job.SettledAt)
|
||||
|
||||
var eventCount int
|
||||
err = tx.QueryRowContext(ctx, `SELECT COUNT(*) FROM batch_image_events WHERE job_id = $1 AND event_type = 'settlement_completed'`, batchID).Scan(&eventCount)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, eventCount)
|
||||
}
|
||||
|
||||
func TestBatchImageRepository_SetBatchImageJobSettlementFailed(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := newBatchImageRepositoryWithSQL(tx)
|
||||
batchID := batchImageTestID(t, "settlement-failed")
|
||||
|
||||
_, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{
|
||||
BatchID: batchID,
|
||||
UserID: 1001,
|
||||
Provider: service.BatchImageProviderGeminiAPI,
|
||||
Model: "gemini-image",
|
||||
Status: service.BatchImageJobStatusSettling,
|
||||
ItemCount: 1,
|
||||
SuccessCount: 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
retryCount, err := repo.SetBatchImageJobSettlementFailed(ctx, batchID, "SETTLEMENT_BILLING_FAILED", "temporary")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, retryCount)
|
||||
|
||||
job, err := repo.GetBatchImageJobByBatchID(ctx, batchID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, service.BatchImageJobStatusSettling, job.Status)
|
||||
require.Equal(t, "SETTLEMENT_BILLING_FAILED", batchImageDerefTest(job.LastErrorCode))
|
||||
require.Equal(t, "temporary", batchImageDerefTest(job.LastErrorMessage))
|
||||
require.Equal(t, 1, job.RetryCount)
|
||||
}
|
||||
|
||||
func TestBatchImageRepository_AppendEvent(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testTx(t)
|
||||
repo := newBatchImageRepositoryWithSQL(tx)
|
||||
batchID := batchImageTestID(t, "event")
|
||||
|
||||
_, err := repo.CreateBatchImageJob(ctx, service.CreateBatchImageJobParams{
|
||||
BatchID: batchID,
|
||||
UserID: 1001,
|
||||
Provider: service.BatchImageProviderVertex,
|
||||
Model: "gemini-2.5-flash-image",
|
||||
ItemCount: 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
err = repo.AppendBatchImageEvent(ctx, batchID, "job_created", map[string]any{"batch_id": batchID})
|
||||
require.NoError(t, err)
|
||||
|
||||
var payload string
|
||||
err = tx.QueryRowContext(ctx, `SELECT payload::text FROM batch_image_events WHERE job_id = $1 AND event_type = 'job_created'`, batchID).Scan(&payload)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, payload, batchID)
|
||||
}
|
||||
|
||||
func batchImageTestID(t *testing.T, prefix string) string {
|
||||
t.Helper()
|
||||
safePrefix := batchImageSafeTestIDSegment(prefix, 20)
|
||||
sum := sha1.Sum([]byte(t.Name()))
|
||||
return "imgbatch_" + safePrefix + "_" + hex.EncodeToString(sum[:])[:16]
|
||||
}
|
||||
|
||||
func batchImageSafeTestIDSegment(v string, maxLen int) string {
|
||||
v = strings.ToLower(strings.TrimSpace(v))
|
||||
v = regexp.MustCompile(`[^a-z0-9_-]+`).ReplaceAllString(v, "-")
|
||||
v = strings.Trim(v, "-_")
|
||||
if v == "" {
|
||||
v = "job"
|
||||
}
|
||||
if len(v) > maxLen {
|
||||
v = v[:maxLen]
|
||||
v = strings.Trim(v, "-_")
|
||||
}
|
||||
if v == "" {
|
||||
return "job"
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func batchImageTestStringPtr(v string) *string {
|
||||
return &v
|
||||
}
|
||||
|
||||
func batchImageDerefTest(v *string) string {
|
||||
if v == nil {
|
||||
return ""
|
||||
}
|
||||
return *v
|
||||
}
|
||||
Reference in New Issue
Block a user