80 lines
2.3 KiB
Go
80 lines
2.3 KiB
Go
|
//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")
|
||
|
|
}
|