82 lines
2.9 KiB
Go
82 lines
2.9 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestWikiProviderCircuitScopeAndExpiry(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
ctx := context.Background()
|
||
|
|
key := wikiProviderCircuitKey(1, t.Name())
|
||
|
|
t.Cleanup(func() { wikiProviderBlocks.Delete(key) })
|
||
|
|
s.blockWikiProvider(ctx, key, errors.New("API request failed with status 402: private provider text"))
|
||
|
|
if err := s.checkWikiProvider(ctx, key); err == nil {
|
||
|
|
t.Fatal("billing failure must block further requests")
|
||
|
|
}
|
||
|
|
if err := s.checkWikiProvider(ctx, wikiProviderCircuitKey(2, t.Name())); err != nil {
|
||
|
|
t.Fatal("different tenant blocked", err)
|
||
|
|
}
|
||
|
|
if err := s.checkWikiProvider(ctx, wikiProviderCircuitKey(1, "another-model")); err != nil {
|
||
|
|
t.Fatal("different model blocked", err)
|
||
|
|
}
|
||
|
|
wikiProviderBlocks.Store(key, wikiProviderBlock{status: 402, until: time.Now().Add(-time.Second)})
|
||
|
|
if err := s.checkWikiProvider(ctx, key); err != nil {
|
||
|
|
t.Fatal("expired block should allow a new attempt", err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiProviderCircuitDoesNotBlockTransientOrOutputErrors(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
key := wikiProviderCircuitKey(1, t.Name())
|
||
|
|
t.Cleanup(func() { wikiProviderBlocks.Delete(key) })
|
||
|
|
for _, message := range []string{"API request failed with status 503", "invalid JSON: 402"} {
|
||
|
|
s.blockWikiProvider(context.Background(), key, errors.New(message))
|
||
|
|
if err := s.checkWikiProvider(context.Background(), key); err != nil {
|
||
|
|
t.Fatalf("unexpected block for %s", message)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiProviderRateLimitUsesShortCooldown(t *testing.T) {
|
||
|
|
s := &wikiIngestService{}
|
||
|
|
key := wikiProviderCircuitKey(1, t.Name())
|
||
|
|
t.Cleanup(func() { wikiProviderBlocks.Delete(key) })
|
||
|
|
s.blockWikiProvider(context.Background(), key, errors.New("API request failed with status 429"))
|
||
|
|
if err := s.checkWikiProvider(context.Background(), key); err == nil || !isLikelyRateLimitError(err) {
|
||
|
|
t.Fatalf("rate limit must be blocked and retain its classification: %v", err)
|
||
|
|
}
|
||
|
|
value, _ := wikiProviderBlocks.Load(key)
|
||
|
|
remaining := time.Until(value.(wikiProviderBlock).until)
|
||
|
|
if remaining <= 0 || remaining > time.Minute {
|
||
|
|
t.Fatalf("unexpected rate limit cooldown: %s", remaining)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestWikiProviderSlotsBoundParallelismAndHonorCancellation(t *testing.T) {
|
||
|
|
key := wikiProviderCircuitKey(1, t.Name())
|
||
|
|
t.Cleanup(func() { wikiProviderSlots.Delete(key) })
|
||
|
|
first, err := acquireWikiProviderSlot(context.Background(), key)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
defer first()
|
||
|
|
second, err := acquireWikiProviderSlot(context.Background(), key)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
defer second()
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
|
||
|
|
defer cancel()
|
||
|
|
if release, err := acquireWikiProviderSlot(ctx, key); release != nil || !errors.Is(err, context.DeadlineExceeded) {
|
||
|
|
t.Fatalf("third call bypassed concurrency limit: %v", err)
|
||
|
|
}
|
||
|
|
other, err := acquireWikiProviderSlot(context.Background(), wikiProviderCircuitKey(2, t.Name()))
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
other()
|
||
|
|
}
|