orchard

  1package main
  2
  3import (
  4	"context"
  5	"os"
  6	"strings"
  7	"sync"
  8	"testing"
  9	"time"
 10)
 11
 12func TestQueueSerialises(t *testing.T) {
 13	q := NewQueue()
 14	var (
 15		mu      sync.Mutex
 16		running int
 17		peak    int
 18		order   []int
 19	)
 20
 21	var wg sync.WaitGroup
 22	for i := 0; i < 5; i++ {
 23		wg.Add(1)
 24		go func(i int) {
 25			defer wg.Done()
 26			release, ok := q.Enter(context.Background(), nil)
 27			if !ok {
 28				t.Errorf("worker %d was refused", i)
 29				return
 30			}
 31			defer release()
 32
 33			mu.Lock()
 34			running++
 35			if running > peak {
 36				peak = running
 37			}
 38			order = append(order, i)
 39			mu.Unlock()
 40
 41			time.Sleep(20 * time.Millisecond)
 42
 43			mu.Lock()
 44			running--
 45			mu.Unlock()
 46		}(i)
 47		// Stagger so the order is deterministic enough to assert on.
 48		time.Sleep(5 * time.Millisecond)
 49	}
 50	wg.Wait()
 51
 52	if peak != 1 {
 53		t.Errorf("peak concurrency was %d, want 1", peak)
 54	}
 55	if len(order) != 5 {
 56		t.Errorf("only %d of 5 ran", len(order))
 57	}
 58}
 59
 60func TestQueueGivesUpOnCancel(t *testing.T) {
 61	q := NewQueue()
 62
 63	held, ok := q.Enter(context.Background(), nil)
 64	if !ok {
 65		t.Fatal("the first caller should run at once")
 66	}
 67
 68	// Second joins and then leaves before its turn.
 69	ctx, cancel := context.WithCancel(context.Background())
 70	left := make(chan bool, 1)
 71	go func() {
 72		_, ok := q.Enter(ctx, nil)
 73		left <- ok
 74	}()
 75	time.Sleep(30 * time.Millisecond)
 76
 77	if n, _ := q.Depth(); n != 1 {
 78		t.Fatalf("depth = %d, want 1 waiting", n)
 79	}
 80	cancel()
 81	if ok := <-left; ok {
 82		t.Error("a cancelled caller should not be given the slot")
 83	}
 84
 85	// A third must still get through once the first releases.
 86	held()
 87	done := make(chan bool, 1)
 88	go func() {
 89		r, ok := q.Enter(context.Background(), nil)
 90		if ok {
 91			r()
 92		}
 93		done <- ok
 94	}()
 95	select {
 96	case ok := <-done:
 97		if !ok {
 98			t.Error("the third caller was refused")
 99		}
100	case <-time.After(2 * time.Second):
101		t.Fatal("the queue stalled after a cancellation")
102	}
103}
104
105func TestQueueReportsPosition(t *testing.T) {
106	q := NewQueue()
107	release, _ := q.Enter(context.Background(), nil)
108
109	seen := make(chan QueueState, 4)
110	go func() {
111		r, ok := q.Enter(context.Background(), func(s QueueState) { seen <- s })
112		if ok {
113			r()
114		}
115	}()
116
117	select {
118	case s := <-seen:
119		if s.Position != 1 || s.Ahead != 0 {
120			t.Errorf("first waiter reported %+v, want position 1 ahead 0", s)
121		}
122	case <-time.After(2 * time.Second):
123		t.Fatal("no position was reported")
124	}
125	release()
126}
127
128// TestLLMConfigMatchesClient guards the two things in llm/config.yaml that the
129// Go side depends on and nothing else would catch: the model name the client
130// sends, and the ttl that makes the model unload at all.
131func TestLLMConfigMatchesClient(t *testing.T) {
132	b, err := os.ReadFile("llm/config.yaml")
133	if err != nil {
134		t.Skip("no llm config here")
135	}
136	cfg := string(b)
137
138	if !strings.Contains(cfg, `"`+NewLLM("", "").Model+`":`) {
139		t.Errorf("llm/config.yaml has no model named %q, which is what the client asks for",
140			NewLLM("", "").Model)
141	}
142	if !strings.Contains(cfg, "ttl:") {
143		t.Error("no ttl in llm/config.yaml, so the model would never unload")
144	}
145}