llmgw/internal/sim/engine_test.go

package sim

import (
	"context"
	"strings"
	"testing"
	"time"
)

// testConfig speeds the simulation up so tests stay fast and deterministic.
func testConfig() Config {
	c := Default()
	c.Tick = time.Millisecond
	c.PrefillTokensPerSec = 1e6
	c.DecodeTokensPerSec = 5000
	c.BatchContention = 0
	c.Seed = 42
	return c
}

// drain collects every token from the channel and waits for completion.
func drain(t *testing.T, ch <-chan Event) (text string, done bool) {
	t.Helper()
	var b strings.Builder
	deadline := time.After(10 * time.Second)
	for {
		select {
		case ev, open := <-ch:
			if !open {
				return b.String(), done
			}
			switch ev.Kind {
			case EventToken:
				b.WriteString(ev.Text)
			case EventDone:
				done = true
			case EventError:
				t.Fatalf("unexpected error: %v", ev.Err)
			}
		case <-deadline:
			t.Fatal("timed out waiting for tokens")
		}
	}
}

func TestEngineGeneratesTokens(t *testing.T) {
	e := New(testConfig())
	defer e.Close()

	ch, err := e.Submit(t.Context(), Prompt{Text: "hello world", MaxTokens: 10})
	if err != nil {
		t.Fatal(err)
	}
	text, done := drain(t, ch)
	if !done {
		t.Error("no EventDone received")
	}
	if got := len(strings.Fields(text)); got != 10 {
		t.Errorf("expected 10 tokens, got %d (%q)", got, text)
	}
}

func TestCancellationReleasesBlocks(t *testing.T) {
	// Essential for testing the router's cancellation propagation: when the
	// client goes away the engine must drop the request and free its KV blocks.
	cfg := testConfig()
	cfg.DecodeTokensPerSec = 20 // slow, so it is still running when cancelled
	e := New(cfg)
	defer e.Close()

	ctx, cancel := context.WithCancel(context.Background())
	ch, err := e.Submit(ctx, Prompt{Text: strings.Repeat("x", 4000), MaxTokens: 1000})
	if err != nil {
		t.Fatal(err)
	}

	waitFor(t, time.Second, func() bool { return e.Stats().Running == 1 })
	if usage := e.Stats().KVCacheUsage; usage <= 0 {
		t.Fatal("running request reserved no KV blocks")
	}

	cancel()

	waitFor(t, 2*time.Second, func() bool {
		s := e.Stats()
		return s.Running == 0 && s.KVCacheUsage == 0
	})

	// The channel must close so the client side never blocks.
	select {
	case _, open := <-ch:
		if open {
			for range ch { //nolint:revive // drain remaining events
			}
		}
	case <-time.After(2 * time.Second):
		t.Fatal("channel stayed open after cancellation")
	}
}

func TestRequestsQueueWhenBatchIsFull(t *testing.T) {
	cfg := testConfig()
	cfg.MaxNumSeqs = 2
	cfg.DecodeTokensPerSec = 20
	e := New(cfg)
	defer e.Close()

	for range 5 {
		if _, err := e.Submit(t.Context(), Prompt{Text: "abc", MaxTokens: 100}); err != nil {
			t.Fatal(err)
		}
	}

	waitFor(t, time.Second, func() bool { return e.Stats().Running == 2 })
	s := e.Stats()
	if s.Running != 2 {
		t.Errorf("running=%d exceeds MaxNumSeqs=2", s.Running)
	}
	if s.Waiting != 3 {
		t.Errorf("waiting=%d, expected 3", s.Waiting)
	}
}

func TestOversizedRequestIsRejected(t *testing.T) {
	cfg := testConfig()
	cfg.NumGPUBlocks = 4
	cfg.BlockSize = 16 // 64 tokens of capacity in total
	e := New(cfg)
	defer e.Close()

	_, err := e.Submit(t.Context(), Prompt{Text: strings.Repeat("x", 10000), MaxTokens: 100})
	if err != ErrTooLong {
		t.Errorf("expected ErrTooLong, got %v", err)
	}
}

func TestPrefixCacheSpeedsUpRepeatedPrompt(t *testing.T) {
	cfg := testConfig()
	cfg.PrefillTokensPerSec = 2000 // make prefill cost measurable
	cfg.DecodeTokensPerSec = 5000
	e := New(cfg)
	defer e.Close()

	prompt := Prompt{Text: strings.Repeat("shared system prompt. ", 200), MaxTokens: 5}

	first := timeToFirstToken(t, e, prompt)
	// Send the same prompt again now that the first one populated the cache.
	second := timeToFirstToken(t, e, prompt)

	if second >= first {
		t.Errorf("prefix cache did not reduce TTFT: first=%v second=%v", first, second)
	}
	if e.Stats().PrefixHits == 0 {
		t.Error("prefix cache hit counter did not advance")
	}
}

func TestNoHitsWhenPrefixCacheDisabled(t *testing.T) {
	cfg := testConfig()
	cfg.PrefixCacheBlocks = 0
	e := New(cfg)
	defer e.Close()

	p := Prompt{Text: strings.Repeat("same prompt. ", 50), MaxTokens: 3}
	for range 3 {
		ch, err := e.Submit(t.Context(), p)
		if err != nil {
			t.Fatal(err)
		}
		drain(t, ch)
	}
	if h := e.Stats().PrefixHits; h != 0 {
		t.Errorf("disabled cache reported %d hits", h)
	}
}

func TestShutdownDrainsPendingRequests(t *testing.T) {
	cfg := testConfig()
	cfg.DecodeTokensPerSec = 1
	e := New(cfg)

	ch, err := e.Submit(t.Context(), Prompt{Text: "abc", MaxTokens: 1000})
	if err != nil {
		t.Fatal(err)
	}
	waitFor(t, time.Second, func() bool { return e.Stats().Running == 1 })
	e.Close()

	deadline := time.After(2 * time.Second)
	for {
		select {
		case _, open := <-ch:
			if !open {
				return
			}
		case <-deadline:
			t.Fatal("channel stayed open after shutdown")
		}
	}
}

func TestInjectedFailure(t *testing.T) {
	cfg := testConfig()
	cfg.FailRate = 1.0
	e := New(cfg)
	defer e.Close()

	if _, err := e.Submit(t.Context(), Prompt{Text: "abc"}); err != ErrInjected {
		t.Errorf("expected ErrInjected, got %v", err)
	}
}

// --- helpers ---

func timeToFirstToken(t *testing.T, e *Engine, p Prompt) time.Duration {
	t.Helper()
	start := time.Now()
	ch, err := e.Submit(t.Context(), p)
	if err != nil {
		t.Fatal(err)
	}
	var ttft time.Duration
	for ev := range ch {
		if ev.Kind == EventToken && ttft == 0 {
			ttft = time.Since(start)
		}
	}
	if ttft == 0 {
		t.Fatal("no tokens were generated")
	}
	return ttft
}

func waitFor(t *testing.T, limit time.Duration, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(limit)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatalf("condition not met within %v", limit)
}