Vectrum
Sign In
106 lines · 2.6 KB · Go
package tui

import (
	"strings"
	"testing"
	"unicode/utf8"

	"github.com/charmbracelet/lipgloss"
	"github.com/muesli/termenv"

	"harness/internal/llm"
)

// forceColour makes styles emit real escape sequences. Without it the test
// renderer detects no terminal, strips every colour, and any assertion about
// what a style actually painted passes vacuously.
func forceColour(t *testing.T) {
	t.Helper()
	prev := lipgloss.ColorProfile()
	lipgloss.SetColorProfile(termenv.TrueColor)
	t.Cleanup(func() { lipgloss.SetColorProfile(prev) })
}

// unbackedCells returns the printable characters that were emitted with no
// background colour in force.
func unbackedCells(s string) string {
	var out []rune
	active := ""
	for i := 0; i < len(s); {
		if s[i] == 0x1b && i+1 < len(s) && s[i+1] == '[' {
			j := i + 2
			for j < len(s) && s[j] != 'm' {
				j++
			}
			if j >= len(s) {
				break
			}
			if p := s[i+2 : j]; p == "" || p == "0" {
				active = ""
			} else {
				active = p
			}
			i = j + 1
			continue
		}
		r, size := utf8.DecodeRuneInString(s[i:])
		if !strings.Contains(active, "48;") {
			out = append(out, r)
		}
		i += size
	}
	return string(out)
}

// A style nested inside a background ends with a reset, which drops the
// background for everything after it and tears a visible hole in the pill.
func TestPillsAreFilledEdgeToEdge(t *testing.T) {
	forceColour(t)

	m := newTestModel(t)
	m.opts.ContextWindow = 200_000
	m.Update(evUsage{usage: llm.Usage{InputTokens: 2_000, CacheReadInputTokens: 18_000}})

	pill, ok := m.contextPill()
	if !ok {
		t.Fatal("no context pill after usage was reported")
	}
	if got := unbackedCells(pill); got != "" {
		t.Errorf("context pill has unfilled cells %q in %q", got, pill)
	}

	if got := unbackedCells(styName.Render(username())); got != "" {
		t.Errorf("name pill has unfilled cells %q", got)
	}
}

func TestContextPillColoursByPressure(t *testing.T) {
	forceColour(t)

	cases := []struct {
		tokens int
		pct    string
		colour lipgloss.AdaptiveColor
	}{
		{10_000, "(10%)", colOK},
		{65_000, "(65%)", colWarn},
		{90_000, "(90%)", colErr},
	}
	for _, tc := range cases {
		m := newTestModel(t)
		m.opts.ContextWindow = 100_000
		m.Update(evUsage{usage: llm.Usage{InputTokens: tc.tokens}})

		pill, ok := m.contextPill()
		if !ok {
			t.Fatalf("%d tokens produced no pill", tc.tokens)
		}
		if !strings.Contains(pill, tc.pct) {
			t.Errorf("%d tokens: %q does not report %s", tc.tokens, pill, tc.pct)
		}
		want := termenv.TrueColor.Color(tc.colour.Dark).Sequence(false)
		if !strings.Contains(pill, want) {
			t.Errorf("%d tokens: pill is not %s: %q", tc.tokens, tc.colour.Dark, pill)
		}
	}
}