Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/scheduler_snapshot_cancellation_test.go

@@ -0,0 +1,79 @@
+//go:build unit
+
+package service
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type schedulerCancellationCache struct {
+	SchedulerCache
+	cancel        context.CancelFunc
+	tokenCaptures int
+}
+
+func (c *schedulerCancellationCache) GetSnapshot(ctx context.Context, _ SchedulerBucket) ([]*Account, bool, error) {
+	c.cancel()
+	return nil, false, ctx.Err()
+}
+
+func (c *schedulerCancellationCache) CaptureBucketWriteToken(ctx context.Context, _ SchedulerBucket) (SchedulerBucketWriteToken, error) {
+	c.tokenCaptures++
+	return SchedulerBucketWriteToken{}, ctx.Err()
+}
+
+func (c *schedulerCancellationCache) GetAccount(ctx context.Context, _ int64) (*Account, error) {
+	c.cancel()
+	return nil, ctx.Err()
+}
+
+type schedulerCancellationAccountRepo struct {
+	AccountRepository
+	listCalls    int
+	getByIDCalls int
+}
+
+func (r *schedulerCancellationAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, _ string) ([]Account, error) {
+	r.listCalls++
+	return nil, ctx.Err()
+}
+
+func (r *schedulerCancellationAccountRepo) GetByID(ctx context.Context, _ int64) (*Account, error) {
+	r.getByIDCalls++
+	return nil, ctx.Err()
+}
+
+func TestSchedulerSnapshotListStopsAfterRequestCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	cache := &schedulerCancellationCache{cancel: cancel}
+	repo := &schedulerCancellationAccountRepo{}
+	svc := NewSchedulerSnapshotService(cache, nil, repo, nil, nil)
+
+	accounts, useMixed, err := svc.ListSchedulableAccounts(ctx, nil, PlatformOpenAI, false)
+
+	require.ErrorIs(t, err, context.Canceled)
+	require.Nil(t, accounts)
+	require.False(t, useMixed)
+	require.Zero(t, cache.tokenCaptures, "canceled requests must not capture a cache publish token")
+	require.Zero(t, repo.listCalls, "canceled requests must not fall back to the database")
+}
+
+func TestSchedulerSnapshotGetAccountStopsAfterRequestCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	cache := &schedulerCancellationCache{cancel: cancel}
+	repo := &schedulerCancellationAccountRepo{}
+	svc := NewSchedulerSnapshotService(cache, nil, repo, nil, nil)
+
+	account, err := svc.GetAccount(ctx, 42)
+
+	require.ErrorIs(t, err, context.Canceled)
+	require.Nil(t, account)
+	require.Zero(t, repo.getByIDCalls, "canceled requests must not fall back to the database")
+}