notarius/internal/framework/llm/scheduled_client_test.go

package llm

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

func TestScheduledClientEnforcesSchedulerLimit(t *testing.T) {
	scheduler, err := NewScheduler(1)
	if err != nil {
		t.Fatalf("NewScheduler() error = %v, want nil", err)
	}
	inner := &blockingStructuredClient{
		release: make(chan struct{}),
	}
	client := NewScheduledClient(inner, scheduler)

	var wg sync.WaitGroup
	for i := 0; i < 3; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			var out map[string]any
			if _, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{}, &out); err != nil {
				t.Errorf("CompleteStructured() error = %v, want nil", err)
			}
		}()
	}

	waitForAtomicAtLeast(t, &inner.calls, 1)
	time.Sleep(20 * time.Millisecond)
	if got := atomic.LoadInt32(&inner.maxInFlight); got > 1 {
		t.Fatalf("max in-flight calls = %d, want <= 1", got)
	}
	close(inner.release)
	wg.Wait()
	if got := atomic.LoadInt32(&inner.calls); got != 3 {
		t.Fatalf("calls = %d, want 3", got)
	}
}

func TestScheduledClientPropagatesClientError(t *testing.T) {
	scheduler, err := NewScheduler(1)
	if err != nil {
		t.Fatalf("NewScheduler() error = %v, want nil", err)
	}
	expected := errors.New("provider unavailable")
	client := NewScheduledClient(&errorStructuredClient{err: expected}, scheduler)

	_, err = client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{}, &struct{}{})
	if !errors.Is(err, expected) {
		t.Fatalf("CompleteStructured() error = %v, want %v", err, expected)
	}
}

func TestScheduledClientPropagatesSchedulerError(t *testing.T) {
	scheduler, err := NewScheduler(1)
	if err != nil {
		t.Fatalf("NewScheduler() error = %v, want nil", err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	client := NewScheduledClient(&errorStructuredClient{}, scheduler)

	_, err = client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{}, &struct{}{})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("CompleteStructured() error = %v, want context canceled", err)
	}
}

func TestScheduledClientPreservesCheckpointFingerprints(t *testing.T) {
	scheduler, err := NewScheduler(1)
	if err != nil {
		t.Fatal(err)
	}
	inner := &fingerprintedStructuredClient{
		fingerprints: []CheckpointFingerprint{
			{Name: "profile_source", Value: "sha256:one"},
			{Name: "backend_target", Value: "sha256:two"},
		},
	}
	client := NewScheduledClient(inner, scheduler)
	provider, ok := client.(CheckpointFingerprintProvider)
	if !ok {
		t.Fatalf("scheduled client %T does not preserve checkpoint fingerprints", client)
	}
	got, err := provider.LLMCheckpointFingerprints()
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != len(inner.fingerprints) ||
		got[0] != inner.fingerprints[0] ||
		got[1] != inner.fingerprints[1] {
		t.Fatalf("checkpoint fingerprints = %#v, want %#v", got, inner.fingerprints)
	}
}

type blockingStructuredClient struct {
	release     chan struct{}
	inFlight    int32
	maxInFlight int32
	calls       int32
}

func (c *blockingStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	atomic.AddInt32(&c.calls, 1)
	current := atomic.AddInt32(&c.inFlight, 1)
	for {
		seen := atomic.LoadInt32(&c.maxInFlight)
		if current <= seen || atomic.CompareAndSwapInt32(&c.maxInFlight, seen, current) {
			break
		}
	}
	defer atomic.AddInt32(&c.inFlight, -1)

	select {
	case <-c.release:
	case <-ctx.Done():
		return contracts.StructuredCompletionResponse{}, ctx.Err()
	}
	if target, ok := out.(*map[string]any); ok {
		*target = map[string]any{"ok": true}
	}
	return contracts.StructuredCompletionResponse{
		Content: json.RawMessage(`{"ok":true}`),
	}, nil
}

type errorStructuredClient struct {
	err error
}

type fingerprintedStructuredClient struct {
	errorStructuredClient
	fingerprints []CheckpointFingerprint
}

func (c *fingerprintedStructuredClient) LLMCheckpointFingerprints() ([]CheckpointFingerprint, error) {
	return append([]CheckpointFingerprint(nil), c.fingerprints...), nil
}

func (c *errorStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	return contracts.StructuredCompletionResponse{}, c.err
}