107 lines
2.8 KiB
Go
107 lines
2.8 KiB
Go
package auth
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
|
|
)
|
|
|
|
func TestHomeDispatchSelectionReleasesAttemptCancelTokensWithoutGrowingResources(t *testing.T) {
|
|
registry := executionregistry.New()
|
|
pending, errBegin := registry.BeginDispatch()
|
|
if errBegin != nil {
|
|
t.Fatal(errBegin)
|
|
}
|
|
scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{})
|
|
if errInstall != nil {
|
|
t.Fatal(errInstall)
|
|
}
|
|
selection, errSelection := newHomeDispatchSelection(&Auth{ID: "home-auth"}, nil, "home", scope)
|
|
if errSelection != nil {
|
|
t.Fatal(errSelection)
|
|
}
|
|
|
|
for range 100 {
|
|
_, release, errAttempt := selection.AttemptContext(context.Background())
|
|
if errAttempt != nil {
|
|
t.Fatalf("AttemptContext() error = %v", errAttempt)
|
|
}
|
|
release()
|
|
}
|
|
|
|
selection.resources.mu.Lock()
|
|
resourceCount := len(selection.resources.closers)
|
|
selection.resources.mu.Unlock()
|
|
if resourceCount != 1 {
|
|
t.Fatalf("bound resources = %d, want 1 attempt cancel registry", resourceCount)
|
|
}
|
|
if got := selection.attemptCancels.Len(); got != 0 {
|
|
t.Fatalf("active attempt cancel tokens = %d, want 0", got)
|
|
}
|
|
|
|
selection.End("completed")
|
|
drainCtx, cancelDrain := context.WithTimeout(context.Background(), time.Second)
|
|
defer cancelDrain()
|
|
if errDrain := registry.Drain(drainCtx); errDrain != nil {
|
|
t.Fatalf("Drain() error = %v", errDrain)
|
|
}
|
|
}
|
|
|
|
func TestAttemptCancelReleaseAfterCloseCancelsOnce(t *testing.T) {
|
|
cancels := &attemptCancels{}
|
|
var cancelCalls atomic.Int32
|
|
release, errAdd := cancels.Add(func() { cancelCalls.Add(1) })
|
|
if errAdd != nil {
|
|
t.Fatalf("Add() error = %v", errAdd)
|
|
}
|
|
if errClose := cancels.Close(); errClose != nil {
|
|
t.Fatalf("Close() error = %v", errClose)
|
|
}
|
|
release()
|
|
if got := cancelCalls.Load(); got != 1 {
|
|
t.Fatalf("cancel calls = %d, want 1", got)
|
|
}
|
|
}
|
|
|
|
func TestHomeDispatchSelectionAttemptReleaseRacesDrainExactlyOnce(t *testing.T) {
|
|
registry := executionregistry.New()
|
|
pending, errBegin := registry.BeginDispatch()
|
|
if errBegin != nil {
|
|
t.Fatal(errBegin)
|
|
}
|
|
scope, errInstall := registry.Install(pending, executionregistry.ScopeSpec{})
|
|
if errInstall != nil {
|
|
t.Fatal(errInstall)
|
|
}
|
|
selection, errSelection := newHomeDispatchSelection(&Auth{ID: "home-auth"}, nil, "home", scope)
|
|
if errSelection != nil {
|
|
t.Fatal(errSelection)
|
|
}
|
|
|
|
_, release, errAttempt := selection.AttemptContext(context.Background())
|
|
if errAttempt != nil {
|
|
t.Fatal(errAttempt)
|
|
}
|
|
var wg sync.WaitGroup
|
|
wg.Add(2)
|
|
go func() {
|
|
defer wg.Done()
|
|
release()
|
|
}()
|
|
go func() {
|
|
defer wg.Done()
|
|
selection.End("draining")
|
|
}()
|
|
wg.Wait()
|
|
|
|
if got := selection.attemptCancels.Len(); got != 0 {
|
|
t.Fatalf("active attempt cancel tokens = %d, want 0", got)
|
|
}
|
|
if errDrain := registry.Drain(context.Background()); errDrain != nil {
|
|
t.Fatalf("Drain() error = %v", errDrain)
|
|
}
|
|
}
|