xinmi/scripts/graph_extraction_probe/budget_test.go

package main

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/models/requestbudget"
)

func TestProbeSendLimit(t *testing.T) {
	one, two, three, zero, four := 1, 2, 3, 0, 4
	for _, tc := range []struct {
		name        string
		requested   *int
		environment string
		envSet      bool
		want        int
	}{
		{name: "default", want: 3},
		{name: "minimum", requested: &one, want: 1},
		{name: "two", requested: &two, environment: "2", envSet: true, want: 2},
		{name: "explicit default", requested: &three, environment: "3", envSet: true, want: 3},
		{name: "default environment", environment: "3", envSet: true, want: 3},
		{name: "zero", requested: &zero},
		{name: "excess", requested: &four},
		{name: "environment mismatch", requested: &two, environment: "3", envSet: true},
		{name: "environment without request", environment: "2", envSet: true},
		{name: "empty environment", requested: &two, envSet: true},
		{name: "padded environment", requested: &two, environment: "02", envSet: true},
		{name: "decimal environment", requested: &two, environment: "2.0", envSet: true},
		{name: "whitespace environment", requested: &two, environment: " 2", envSet: true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			got, err := probeSendLimit(tc.requested, tc.environment, tc.envSet)
			if got != tc.want || (err != nil) != (tc.want == 0) {
				t.Fatalf("limit = %d, error = %v; want %d", got, err, tc.want)
			}
		})
	}
}

type probeRoundTripper func(*http.Request) (*http.Response, error)

func (f probeRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) }

func TestProbeTransportHonorsActualSendCap(t *testing.T) {
	for _, limit := range []int{1, 2, requestbudget.GraphMaxSends} {
		t.Run(fmt.Sprint(limit), func(t *testing.T) {
			ctx, guard, cancel, err := bindProbeExecution(context.Background(), "offline", "model", limit)
			if err != nil {
				t.Fatal(err)
			}
			defer cancel()
			ctx = requestbudget.WithModel(ctx, "model")
			sends := 0
			transport := requestbudget.Transport{Base: probeRoundTripper(func(req *http.Request) (*http.Response, error) {
				sends++
				if req.Context().Err() != nil {
					t.Fatal("allowed response was canceled")
				}
				return &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("usable response"))}, nil
			})}
			for i := 0; i < limit; i++ {
				req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "https://offline.invalid", nil)
				response, err := transport.RoundTrip(req)
				if err != nil {
					t.Fatalf("send %d: %v", i+1, err)
				}
				body, readErr := io.ReadAll(response.Body)
				response.Body.Close()
				if readErr != nil || string(body) != "usable response" || requestbudget.Check(ctx) != nil {
					t.Fatal("last allowed response must remain usable")
				}
			}
			req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "https://offline.invalid", nil)
			if _, err := transport.RoundTrip(req); !errors.Is(err, requestbudget.ErrStopped) {
				t.Fatalf("excess send error = %v", err)
			}
			if sends != limit || guard.Reservations() != limit || !errors.Is(context.Cause(ctx), requestbudget.ErrStopped) {
				t.Fatalf("sends = %d, reservations = %d, cause = %v", sends, guard.Reservations(), context.Cause(ctx))
			}
			if err := guard.Reserve(context.Background(), "model", requestbudget.GraphPhase); !errors.Is(err, requestbudget.ErrStopped) {
				t.Fatal("stopped execution was revived")
			}
		})
	}
}

func TestProbeBudgetPreservesScopeAndDeadline(t *testing.T) {
	for _, tc := range []struct{ name, model, phase string }{
		{"model", "other", requestbudget.GraphPhase},
		{"phase", "model", "other"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			ctx, guard, cancel, err := bindProbeExecution(context.Background(), "offline", "model", 2)
			if err != nil {
				t.Fatal(err)
			}
			defer cancel()
			if err := guard.Reserve(ctx, tc.model, tc.phase); !errors.Is(err, requestbudget.ErrStopped) {
				t.Fatalf("scope mismatch error = %v", err)
			}
			if guard.Reservations() != 0 || !errors.Is(context.Cause(ctx), requestbudget.ErrStopped) {
				t.Fatal("scope mismatch consumed a send or did not stop execution")
			}
		})
	}
	t.Run("inherited deadline", func(t *testing.T) {
		deadline := time.Now().Add(time.Minute)
		parent, parentCancel := context.WithDeadline(context.Background(), deadline)
		defer parentCancel()
		ctx, guard, cancel, err := bindProbeExecution(parent, "offline", "model", 2)
		if err != nil {
			t.Fatal(err)
		}
		defer cancel()
		if got, ok := ctx.Deadline(); !ok || !got.Equal(deadline) {
			t.Fatal("inherited deadline was reset")
		}
		if err := guard.Reserve(ctx, "model", requestbudget.GraphPhase); err != nil {
			t.Fatal(err)
		}
		if got, _ := ctx.Deadline(); !got.Equal(deadline) {
			t.Fatal("reservation reset the deadline")
		}
		parentCancel()
		if err := guard.Reserve(ctx, "model", requestbudget.GraphPhase); !errors.Is(err, context.Canceled) {
			t.Fatalf("cancellation error = %v", err)
		}
		if guard.Reservations() != 1 {
			t.Fatal("canceled send consumed a reservation")
		}
	})
	t.Run("default deadline", func(t *testing.T) {
		started := time.Now()
		ctx, _, cancel, err := bindProbeExecution(context.Background(), "offline", "model", 2)
		if err != nil {
			t.Fatal(err)
		}
		defer cancel()
		deadline, ok := ctx.Deadline()
		if !ok || deadline.Before(started.Add(requestbudget.GraphTimeout)) || deadline.After(time.Now().Add(requestbudget.GraphTimeout)) {
			t.Fatal("probe did not retain the shared absolute timeout")
		}
	})
	t.Run("already canceled", func(t *testing.T) {
		parent, cancel := context.WithCancel(context.Background())
		cancel()
		if _, _, cleanup, err := bindProbeExecution(parent, "offline", "model", 2); err == nil {
			cleanup()
			t.Fatal("canceled execution was rebound")
		}
	})
	t.Run("nested execution", func(t *testing.T) {
		ctx, _, cancel, err := bindProbeExecution(context.Background(), "offline", "model", 2)
		if err != nil {
			t.Fatal(err)
		}
		defer cancel()
		if _, _, cleanup, err := bindProbeExecution(ctx, "retry", "model", 2); err == nil {
			cleanup()
			t.Fatal("nested execution reset the allowance")
		}
	})
}

func TestProbeBudgetConcurrentReservationsStayBounded(t *testing.T) {
	ctx, guard, cancel, err := bindProbeExecution(context.Background(), "offline", "model", 2)
	if err != nil {
		t.Fatal(err)
	}
	defer cancel()
	var successful atomic.Int32
	var group sync.WaitGroup
	for i := 0; i < 16; i++ {
		group.Add(1)
		go func() {
			defer group.Done()
			if guard.Reserve(ctx, "model", requestbudget.GraphPhase) == nil {
				successful.Add(1)
			}
		}()
	}
	group.Wait()
	if successful.Load() != 2 || guard.Reservations() != 2 {
		t.Fatalf("successful = %d; reservations = %d", successful.Load(), guard.Reservations())
	}
}