51 lines
983 B
Go
51 lines
983 B
Go
package agent_test
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"chat-agent-system/internal/agent"
|
|
)
|
|
|
|
type slowBackend struct {
|
|
running atomic.Int32
|
|
peak atomic.Int32
|
|
}
|
|
|
|
func (s *slowBackend) Name() string { return "slow" }
|
|
|
|
func (s *slowBackend) Execute(ctx context.Context, req agent.Request) (*agent.Result, error) {
|
|
cur := s.running.Add(1)
|
|
for {
|
|
p := s.peak.Load()
|
|
if cur <= p || s.peak.CompareAndSwap(p, cur) {
|
|
break
|
|
}
|
|
}
|
|
time.Sleep(40 * time.Millisecond)
|
|
s.running.Add(-1)
|
|
return &agent.Result{Success: true}, nil
|
|
}
|
|
|
|
func TestLimitingBackendCapsConcurrency(t *testing.T) {
|
|
inner := &slowBackend{}
|
|
limited := agent.WrapWithConcurrencyLimit(inner, 2)
|
|
|
|
var wg sync.WaitGroup
|
|
for i := 0; i < 6; i++ {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
_, _ = limited.Execute(context.Background(), agent.Request{})
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
|
|
if peak := inner.peak.Load(); peak > 2 {
|
|
t.Fatalf("expected peak concurrency <= 2, got %d", peak)
|
|
}
|
|
}
|