99 lines
2.8 KiB
Go
99 lines
2.8 KiB
Go
|
package requestbudget
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"sync/atomic"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestGraphExecutionRefusesFourthSendAndScopeChanges(t *testing.T) {
|
||
|
|
for _, mismatch := range []string{"limit", "model", "phase"} {
|
||
|
|
t.Run(mismatch, func(t *testing.T) {
|
||
|
|
ctx, guard, cancel, err := BindGraphExecution(context.Background(), "spent-receipt", "model")
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
defer cancel()
|
||
|
|
ctx = WithModel(ctx, "model")
|
||
|
|
if mismatch == "model" {
|
||
|
|
ctx = WithModel(ctx, "changed")
|
||
|
|
}
|
||
|
|
if mismatch == "phase" {
|
||
|
|
ctx = WithPhase(ctx, "changed")
|
||
|
|
}
|
||
|
|
sends := 0
|
||
|
|
transport := Transport{Base: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||
|
|
sends++
|
||
|
|
return nil, errors.New("provider failure")
|
||
|
|
})}
|
||
|
|
for i := 0; i < 4; i++ {
|
||
|
|
req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://local.test", nil)
|
||
|
|
_, _ = transport.RoundTrip(req)
|
||
|
|
}
|
||
|
|
want := 0
|
||
|
|
if mismatch == "limit" {
|
||
|
|
want = 3
|
||
|
|
}
|
||
|
|
if sends != want || guard.Reservations() != want || !errors.Is(guard.Check(ctx), ErrStopped) {
|
||
|
|
t.Fatalf("sends=%d reserved=%d stop=%v", sends, guard.Reservations(), guard.Check(ctx))
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestGraphExecutionCancelsBlockedHTTPAtInheritedDeadline(t *testing.T) {
|
||
|
|
var sends atomic.Int32
|
||
|
|
arrived := make(chan struct{})
|
||
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
|
|
sends.Add(1)
|
||
|
|
close(arrived)
|
||
|
|
<-r.Context().Done()
|
||
|
|
}))
|
||
|
|
defer server.Close()
|
||
|
|
parent, parentCancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
|
||
|
|
defer parentCancel()
|
||
|
|
ctx, guard, cancel, err := BindGraphExecution(parent, "spent-receipt", "model")
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
defer cancel()
|
||
|
|
wantDeadline, _ := parent.Deadline()
|
||
|
|
if deadline, _ := ctx.Deadline(); !deadline.Equal(wantDeadline) {
|
||
|
|
t.Fatal("deadline was replenished")
|
||
|
|
}
|
||
|
|
req, _ := http.NewRequestWithContext(WithModel(ctx, "model"), http.MethodGet, server.URL, nil)
|
||
|
|
started := time.Now()
|
||
|
|
_, err = (&http.Client{Transport: Transport{}}).Do(req)
|
||
|
|
if !errors.Is(err, context.DeadlineExceeded) || time.Since(started) > 2*time.Second {
|
||
|
|
t.Fatalf("blocked send did not cancel at deadline: %v", err)
|
||
|
|
}
|
||
|
|
select {
|
||
|
|
case <-arrived:
|
||
|
|
default:
|
||
|
|
t.Fatal("request never reached local server")
|
||
|
|
}
|
||
|
|
if sends.Load() != 1 || guard.Reservations() != 1 {
|
||
|
|
t.Fatal("blocked send was not counted")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestGraphExecutionLastAllowedResponseRemainsUsable(t *testing.T) {
|
||
|
|
ctx, guard, cancel, err := BindGraphExecution(context.Background(), "spent-receipt", "model")
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
defer cancel()
|
||
|
|
for i := 0; i < GraphMaxSends; i++ {
|
||
|
|
if err := guard.Reserve(ctx, "model", GraphPhase); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if err := guard.Check(ctx); err != nil {
|
||
|
|
t.Fatalf("third valid response rejected: %v", err)
|
||
|
|
}
|
||
|
|
}
|