58 lines
1.8 KiB
Go
58 lines
1.8 KiB
Go
package queue
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"git.ilapage.cn/OPC/chorus/internal/core/model"
|
|
)
|
|
|
|
type fakeRepository struct {
|
|
claims int
|
|
}
|
|
|
|
func (f *fakeRepository) CreateIdempotent(context.Context, *model.Generation, []model.GenerationInput) (bool, error) {
|
|
return false, nil
|
|
}
|
|
func (f *fakeRepository) ClaimNext(context.Context, string, time.Duration) (*Claim, error) {
|
|
f.claims++
|
|
return nil, nil
|
|
}
|
|
func (f *fakeRepository) BeginProviderAttempt(context.Context, uint64, string, uint64, uint64) (model.Attempt, bool, error) {
|
|
return model.Attempt{}, false, nil
|
|
}
|
|
func (f *fakeRepository) FinishProviderAttempt(context.Context, uint64, string, model.Attempt) (bool, error) {
|
|
return false, nil
|
|
}
|
|
func (f *fakeRepository) Defer(context.Context, uint64, string, time.Time) (bool, error) {
|
|
return false, nil
|
|
}
|
|
func (f *fakeRepository) AssignProvider(context.Context, uint64, string, uint64) (bool, error) {
|
|
return false, nil
|
|
}
|
|
func (f *fakeRepository) Succeed(context.Context, uint64, string, []model.GenerationOutput, model.Attempt) (bool, error) {
|
|
return false, nil
|
|
}
|
|
func (f *fakeRepository) Fail(context.Context, uint64, string, string, string, model.Attempt) (bool, error) {
|
|
return false, nil
|
|
}
|
|
func (f *fakeRepository) Inputs(context.Context, uint64) ([]model.GenerationInput, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func TestControllerStopsNewClaims(t *testing.T) {
|
|
repository := &fakeRepository{}
|
|
controller := NewController(repository)
|
|
if _, err := controller.ClaimNext(context.Background(), "worker", time.Second); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
controller.StopClaims()
|
|
if _, err := controller.ClaimNext(context.Background(), "worker", time.Second); err != ErrClaimsStopped {
|
|
t.Fatalf("ClaimNext() error = %v", err)
|
|
}
|
|
if repository.claims != 1 {
|
|
t.Fatalf("repository claims = %d", repository.claims)
|
|
}
|
|
}
|