orchard

 1package main
 2
 3import (
 4	"strings"
 5	"testing"
 6)
 7
 8// The exact shape Ornith leaked into an answer: no closing function tag, no
 9// closing parameter tag, and the whole thing rendered to the reader as markup.
10const leaked = "\n<tool_call> <function=web_fetch> <parameter=url> https://github.com/restic/restic/releases </tool_call>\n"
11
12func known(n string) bool { return n == "web_fetch" || n == "web_search" }
13
14func TestSalvageRecoversTheCall(t *testing.T) {
15	clean, calls := salvageCalls(leaked, known)
16	if len(calls) != 1 {
17		t.Fatalf("want 1 call, got %d", len(calls))
18	}
19	if calls[0].Function.Name != "web_fetch" {
20		t.Errorf("name = %q", calls[0].Function.Name)
21	}
22	if !strings.Contains(calls[0].Function.Arguments, "restic/restic/releases") {
23		t.Errorf("args = %q", calls[0].Function.Arguments)
24	}
25	if strings.Contains(clean, "<") {
26		t.Errorf("markup survived into the answer: %q", clean)
27	}
28}
29
30func TestSalvageStripsMarkupEvenWhenTheToolIsUnknown(t *testing.T) {
31	clean, calls := salvageCalls("Here you go.\n"+leaked, func(string) bool { return false })
32	if len(calls) != 0 {
33		t.Fatalf("want no calls for an unknown tool, got %d", len(calls))
34	}
35	if strings.Contains(clean, "tool_call") || strings.Contains(clean, "<") {
36		t.Errorf("markup survived: %q", clean)
37	}
38	if !strings.Contains(clean, "Here you go.") {
39		t.Errorf("real prose was lost: %q", clean)
40	}
41}
42
43func TestSalvageJSONShape(t *testing.T) {
44	in := `<tool_call>{"name": "web_search", "arguments": {"query": "restic latest release"}}</tool_call>`
45	clean, calls := salvageCalls(in, known)
46	if len(calls) != 1 || calls[0].Function.Name != "web_search" {
47		t.Fatalf("calls = %+v", calls)
48	}
49	if !strings.Contains(calls[0].Function.Arguments, "restic latest release") {
50		t.Errorf("args = %q", calls[0].Function.Arguments)
51	}
52	if clean != "" {
53		t.Errorf("clean = %q, want empty", clean)
54	}
55}
56
57func TestSalvageLeavesOrdinaryProseAlone(t *testing.T) {
58	in := "Use `restic backup` nightly. Compare a < b in your script."
59	clean, calls := salvageCalls(in, known)
60	if len(calls) != 0 || clean != in {
61		t.Errorf("prose was modified: %q, calls %d", clean, len(calls))
62	}
63}