Vectrum
Sign In
72 lines · 2.3 KB · Go
package tools

import (
	"strings"
	"testing"
)

func TestContextStatusReportsTheNumbersThatMatter(t *testing.T) {
	tool := NewContextStatus(func() ContextReport {
		return ContextReport{
			UsedTokens: 62_000, WindowTokens: 1_000_000, CompactAt: 0.7,
			Turns: 50, CostUSD: 2.1, CostKnown: true,
		}
	})
	got, err := run(t, tool, struct{}{})
	if err != nil {
		t.Fatalf("context_status: %v", err)
	}
	for _, want := range []string{"62,000", "1,000,000", "6%", "938,000", "700,000", "Turns so far: 50", "$2.10"} {
		if !strings.Contains(got, want) {
			t.Errorf("report does not mention %q:\n%s", want, got)
		}
	}
}

// A model whose pricing Codey does not know gets no figure rather than a
// wrong one.
func TestContextStatusOmitsUnknownCost(t *testing.T) {
	got, _ := run(t, NewContextStatus(func() ContextReport {
		return ContextReport{UsedTokens: 10, WindowTokens: 100, Turns: 1}
	}), struct{}{})
	if strings.Contains(got, "$") {
		t.Errorf("an unknown price must not be printed at all:\n%s", got)
	}
}

// Past the threshold the answer has to say the fold is imminent, because that
// is the point at which the model should wrap up rather than start something.
func TestContextStatusWarnsWhenTheFoldIsDue(t *testing.T) {
	got, _ := run(t, NewContextStatus(func() ContextReport {
		return ContextReport{UsedTokens: 800, WindowTokens: 1000, CompactAt: 0.7, Turns: 9}
	}), struct{}{})
	if !strings.Contains(got, "fold") || !strings.Contains(got, "summary") {
		t.Errorf("no warning that compaction is next:\n%s", got)
	}
}

func TestContextStatusSurvivesAnUnknownWindow(t *testing.T) {
	got, err := run(t, NewContextStatus(func() ContextReport {
		return ContextReport{UsedTokens: 1234, Turns: 2}
	}), struct{}{})
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(got, "1,234") || !strings.Contains(got, "not known") {
		t.Errorf("an unknown window should be said plainly:\n%s", got)
	}
}

func TestContextStatusIsNotRegisteredWithoutASource(t *testing.T) {
	if NewContextStatus(nil) != nil {
		t.Error("with no way to read the numbers the tool must not exist")
	}
}

func TestThousands(t *testing.T) {
	for in, want := range map[int]string{0: "0", 999: "999", 1000: "1,000", 970000: "970,000", 1234567: "1,234,567", -5: "0"} {
		if got := thousands(in); got != want {
			t.Errorf("thousands(%d) = %q, want %q", in, got, want)
		}
	}
}