audita/examples/library-adoption/capacity_test.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"sync"
	"testing"
	"time"

	ck "gitea.maximumdirect.net/eric/ConfigKit"
	pk "gitea.maximumdirect.net/eric/promptkit"
)

func roleVars() map[string]string {
	return map[string]string{"SectionJSON": "[]", "GlossaryJSON": "[]", "PayloadJSON": "[]", "TranscriptDescriptionBlock": ""}
}
func backendOverride(field string, n int) ck.Override {
	return ck.Override{Path: ck.Keys("backends", "offline", field), Value: n, Flag: "--" + field}
}
func receive[T any](t *testing.T, ch <-chan T) T {
	t.Helper()
	select {
	case v := <-ch:
		return v
	case <-time.After(5 * time.Second):
		t.Fatal("controlled operation did not finish")
		var zero T
		return zero
	}
}

func TestSharedCapacityAndRecovery(t *testing.T) {
	for _, limit := range []int{4, 2} {
		t.Run(map[int]string{4: "local", 2: "api"}[limit], func(t *testing.T) {
			var mu sync.Mutex
			active, peak, calls := 0, 0, 0
			entered := make(chan struct{}, limit)
			release := make(chan struct{})
			var releaseOnce sync.Once
			unblock := func() { releaseOnce.Do(func() { close(release) }) }
			defer unblock()
			transport := roundTripFunc(func(req *http.Request) (*http.Response, error) {
				mu.Lock()
				active++
				calls++
				n := calls
				if active > peak {
					peak = active
				}
				mu.Unlock()
				defer func() { mu.Lock(); active--; mu.Unlock() }()
				if n <= limit {
					entered <- struct{}{}
					select {
					case <-release:
					case <-req.Context().Done():
						return nil, req.Context().Err()
					}
				}
				var body struct {
					Model string `json:"model"`
				}
				if err := json.NewDecoder(req.Body).Decode(&body); err != nil {
					return nil, err
				}
				if n == 1 {
					return providerResponse(503, "", nil), nil
				}
				if n == 2 {
					return providerResponse(200, "invalid JSON", nil), nil
				}
				if body.Model == "synthetic-grammar" {
					return providerResponse(200, validProposal, nil), nil
				}
				return providerResponse(200, validAssessment, nil), nil
			})
			r := testRuntime(t, transport, backendOverride("concurrency", limit), backendOverride("queue_capacity", 0))
			recorder := &Recorder{}
			results := make(chan error, limit)
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			for i := 0; i < limit; i++ {
				role := []string{"proposal", "assessment"}[i%2]
				go func() {
					result, err := r.call(ctx, role, roleVars(), recorder.observer(role))
					if err == nil && !result.Validation.IsValid {
						err = errors.New("unexpected invalid output")
					}
					results <- err
				}()
			}
			for i := 0; i < limit; i++ {
				receive(t, entered)
			}
			_, err := r.call(ctx, "assessment", roleVars(), recorder.observer("assessment"))
			if !errors.Is(err, pk.ErrCapacityExceeded) {
				t.Fatalf("full queue: %v", err)
			}
			unblock()
			for i := 0; i < limit; i++ {
				if err := receive(t, results); err != nil {
					t.Fatal(err)
				}
			}
			// Rejected admission and earlier retry/repair must leave capacity usable.
			if _, err := r.call(ctx, "proposal", roleVars(), recorder.observer("proposal")); err != nil {
				t.Fatal(err)
			}
			mu.Lock()
			observedPeak := peak
			mu.Unlock()
			if observedPeak != limit {
				t.Fatalf("peak %d, limit %d", observedPeak, limit)
			}
			_, retries, repairs := recorder.totals()
			if retries != 1 || repairs != 1 {
				t.Fatalf("retry/repair activity missing: %d/%d", retries, repairs)
			}
		})
	}
}

func TestQueuedCancellationAndDeadline(t *testing.T) {
	entered := make(chan struct{}, 1)
	release := make(chan struct{})
	var once sync.Once
	defer once.Do(func() { close(release) })
	r := testRuntime(t, roundTripFunc(func(req *http.Request) (*http.Response, error) {
		select {
		case entered <- struct{}{}:
		default:
		}
		select {
		case <-release:
			return providerResponse(200, validProposal, nil), nil
		case <-req.Context().Done():
			return nil, req.Context().Err()
		}
	}), backendOverride("concurrency", 1), backendOverride("queue_capacity", 1))
	recorder := &Recorder{}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	first := make(chan error, 1)
	go func() { _, err := r.call(ctx, "proposal", roleVars(), recorder.observer("proposal")); first <- err }()
	receive(t, entered)
	for _, deadline := range []bool{false, true} {
		pending, cancelPending := context.WithCancel(ctx)
		if deadline {
			cancelPending()
			pending, cancelPending = context.WithTimeout(ctx, 20*time.Millisecond)
		}
		started := make(chan struct{})
		done := make(chan error, 1)
		go func() {
			_, err := r.call(pending, "assessment", roleVars(), func(e pk.ExecutionEvent) {
				recorder.observer("assessment")(e)
				if e.Kind == pk.ExecutionStarted {
					close(started)
				}
			})
			done <- err
		}()
		receive(t, started)
		if !deadline {
			cancelPending()
		}
		err := receive(t, done)
		cancelPending()
		expected := context.Canceled
		if deadline {
			expected = context.DeadlineExceeded
		}
		if !errors.Is(err, expected) {
			t.Fatalf("pending call: %v", err)
		}
	}
	once.Do(func() { close(release) })
	if err := receive(t, first); err != nil {
		t.Fatal(err)
	}
	if _, err := r.call(ctx, "proposal", roleVars(), recorder.observer("proposal")); err != nil {
		t.Fatal(err)
	}
	for _, o := range recorder.snapshot() {
		if o.Role == "assessment" && o.Kind == string(pk.ExecutionAttempt) && o.Dispatched {
			t.Fatal("queued canceled call dispatched")
		}
	}
}

func TestDeadlineStopsRetryAfter(t *testing.T) {
	r := testRuntime(t, roundTripFunc(func(req *http.Request) (*http.Response, error) {
		resp := providerResponse(429, "", nil)
		resp.Header.Set("Retry-After", "60")
		return resp, nil
	}))
	r.Config.LLM.OperationTimeout = time.Second
	recorder := &Recorder{}
	_, err := r.call(context.Background(), "proposal", roleVars(), recorder.observer("proposal"))
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatal(err)
	}
	attempts, _, _ := recorder.totals()
	if attempts != 1 {
		t.Fatalf("attempts %d", attempts)
	}
	found := false
	for _, o := range recorder.snapshot() {
		if o.Kind == string(pk.ExecutionRetryDelay) {
			found = true
		}
	}
	if !found {
		t.Fatal("retry wait not observed")
	}
}