orchard

  1package main
  2
  3import (
  4	"path/filepath"
  5	"strings"
  6	"testing"
  7)
  8
  9// A step goes out as it happens, so the list fills in while the answer is still
 10// being written rather than appearing whole at the end.
 11func TestAStepIsSentAsItHappens(t *testing.T) {
 12	var sent []Event
 13	tr := NewTrace(func(e Event) { sent = append(sent, e) })
 14
 15	tr.Add(Step{Kind: "tool", Label: "wikipedia"})
 16	tr.Add(Step{Kind: "answer", Label: "wrote the reply"})
 17
 18	if len(sent) != 2 {
 19		t.Fatalf("%d events sent, want 2", len(sent))
 20	}
 21	if sent[0].Kind != "step" || sent[0].Step == nil || sent[0].Step.Label != "wikipedia" {
 22		t.Errorf("first event is not the step: %+v", sent[0])
 23	}
 24	if len(tr.Steps()) != 2 {
 25		t.Errorf("%d steps kept, want 2", len(tr.Steps()))
 26	}
 27}
 28
 29// The trace is written next to every message, so a prompt of several thousand
 30// characters cannot go in whole.
 31func TestALongFieldIsClippedAndSaysSo(t *testing.T) {
 32	tr := NewTrace(nil)
 33	tr.Add(Step{Kind: "prompt", In: strings.Repeat("x", stepMax+500)})
 34
 35	got := tr.Steps()[0].In
 36	if len(got) > stepMax+80 {
 37		t.Errorf("field is %d characters, want it clipped near %d", len(got), stepMax)
 38	}
 39	if !strings.Contains(got, "more characters") {
 40		t.Errorf("a clipped field does not say it was clipped: %q", got[len(got)-40:])
 41	}
 42}
 43
 44// A turn that goes wrong must not write a megabyte of steps.
 45func TestTheStepCountIsCapped(t *testing.T) {
 46	tr := NewTrace(nil)
 47	for i := 0; i < stepsMax+25; i++ {
 48		tr.Add(Step{Kind: "model", Label: "round"})
 49	}
 50	if n := len(tr.Steps()); n != stepsMax {
 51		t.Errorf("%d steps kept, want the cap of %d", n, stepsMax)
 52	}
 53}
 54
 55// A nil trace is what the tests and any caller that does not want one pass, and
 56// it has to be usable rather than a panic waiting in the turn loop.
 57func TestANilTraceIsSafe(t *testing.T) {
 58	var tr *Trace
 59	tr.Add(Step{Kind: "model"})
 60	if tr.Steps() != nil {
 61		t.Error("a nil trace returned steps")
 62	}
 63}
 64
 65// The steps have to survive a reload, which means the column has to exist and
 66// be read back. A new column that went into one of the two table creates and
 67// not the other would lose them silently.
 68func TestStepsSurviveAStoreRoundTrip(t *testing.T) {
 69	store, err := OpenStore(filepath.Join(t.TempDir(), "chat.db"))
 70	if err != nil {
 71		t.Fatalf("open: %v", err)
 72	}
 73	defer store.Close()
 74
 75	id, err := store.NewConversation("test")
 76	if err != nil {
 77		t.Fatalf("new conversation: %v", err)
 78	}
 79	want := []Step{
 80		{Kind: "wikipedia", Label: "looked the subject up", In: "north korea", Out: "a country", MS: 18},
 81		{Kind: "gate", Label: "checked the draft", Bad: true},
 82	}
 83	if err := store.Append(id, Stored{Role: RoleAssistant, Content: "hi", Steps: want}); err != nil {
 84		t.Fatalf("append: %v", err)
 85	}
 86
 87	msgs, err := store.Messages(id)
 88	if err != nil {
 89		t.Fatalf("messages: %v", err)
 90	}
 91	if len(msgs) != 1 {
 92		t.Fatalf("%d messages, want 1", len(msgs))
 93	}
 94	got := msgs[0].Steps
 95	if len(got) != 2 {
 96		t.Fatalf("%d steps came back, want 2", len(got))
 97	}
 98	if got[0].In != "north korea" || got[0].MS != 18 {
 99		t.Errorf("the first step came back wrong: %+v", got[0])
100	}
101	if !got[1].Bad {
102		t.Error("the failed step came back as fine")
103	}
104}