Vectrum
Sign In
410 lines · 11.9 KB · Go
package tui

import (
	"bytes"
	"errors"
	"image"
	"image/color"
	"image/png"
	"os"
	"path/filepath"
	"strings"
	"testing"

	tea "github.com/charmbracelet/bubbletea"

	"github.com/johannSo/codey/internal/attach"
	"github.com/johannSo/codey/internal/llm"
)

func testPNG(t *testing.T) []byte {
	t.Helper()
	im := image.NewRGBA(image.Rect(0, 0, 1, 1))
	im.Set(0, 0, color.RGBA{R: 9, G: 9, B: 9, A: 255})
	var b bytes.Buffer
	if err := png.Encode(&b, im); err != nil {
		t.Fatal(err)
	}
	return b.Bytes()
}

// pasteModel is a model whose clipboard holds whatever the test says it does,
// with no desktop anywhere near it.
func pasteModel(t *testing.T, im attach.Image, err error) *model {
	t.Helper()
	m := newTestModel(t)
	m.grabFn = func() (attach.Image, error) { return im, err }
	// No test reaches for the real clipboard: an empty one is the default,
	// and a test that wants words on it says so.
	m.textFn = func() (string, error) { return "", nil }
	return m
}

func clipImage(t *testing.T) attach.Image {
	t.Helper()
	im, err := attach.New(testPNG(t), "clipboard")
	if err != nil {
		t.Fatal(err)
	}
	return im
}

// run drives a command the way the runtime would: call it, feed what comes
// back to Update.
func run(t *testing.T, m *model, cmd tea.Cmd) {
	t.Helper()
	for cmd != nil {
		msg := cmd()
		if msg == nil {
			return
		}
		_, cmd = m.Update(msg)
	}
}

func pasteKey(text string) tea.KeyMsg {
	return tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune(text), Paste: true}
}

func TestCtrlVPutsAMarkerInThePrompt(t *testing.T) {
	m := pasteModel(t, clipImage(t), nil)

	_, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlV})
	run(t, m, cmd)

	if got := m.input.Value(); !strings.Contains(got, "[image 1]") {
		t.Errorf("prompt = %q; want it to hold [image 1]", got)
	}
	if m.imgs.Len() != 1 {
		t.Errorf("%d images attached, want 1", m.imgs.Len())
	}
}

// The marker has to keep clear of the words already typed, or it stops
// reading as one thing.
func TestAMarkerDoesNotRunIntoTheWordBeforeIt(t *testing.T) {
	m := pasteModel(t, clipImage(t), nil)
	m.input.SetValue("look at this")

	_, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlV})
	run(t, m, cmd)

	if got := m.input.Value(); !strings.Contains(got, "this [image 1]") {
		t.Errorf("prompt = %q", got)
	}
}

func TestTwoPastesAreNumberedInOrder(t *testing.T) {
	m := pasteModel(t, clipImage(t), nil)
	for i := 0; i < 2; i++ {
		_, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlV})
		run(t, m, cmd)
	}
	got := m.input.Value()
	if !strings.Contains(got, "[image 1]") || !strings.Contains(got, "[image 2]") {
		t.Errorf("prompt = %q; want both markers", got)
	}
}

// ctrl-V on a clipboard holding text is the common case. It must not put a
// complaint in the transcript every time.
func TestCtrlVWithNoImageSaysNothing(t *testing.T) {
	m := pasteModel(t, attach.Image{}, attach.ErrNoImage)
	before := len(m.tr.entries)

	_, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlV})
	run(t, m, cmd)

	if m.imgs.Len() != 0 {
		t.Errorf("%d images attached from an empty clipboard", m.imgs.Len())
	}
	if len(m.tr.entries) != before {
		t.Errorf("the transcript grew by %d lines", len(m.tr.entries)-before)
	}
}

// A terminal that eats ctrl-V pastes the clipboard's text form, which for a
// screenshot is nothing. That empty paste is the only signal there is, so it
// is taken as one.
func TestAnEmptyPasteReachesForTheClipboard(t *testing.T) {
	m := pasteModel(t, clipImage(t), nil)

	_, cmd := m.Update(pasteKey(""))
	run(t, m, cmd)

	if !strings.Contains(m.input.Value(), "[image 1]") {
		t.Errorf("prompt = %q", m.input.Value())
	}
}

// Dropping a file on the window pastes its path.
func TestPastingAPathAttachesTheFile(t *testing.T) {
	p := filepath.Join(t.TempDir(), "screenshot.png")
	if err := os.WriteFile(p, testPNG(t), 0o644); err != nil {
		t.Fatal(err)
	}
	m := pasteModel(t, attach.Image{}, attach.ErrNoImage)

	_, cmd := m.Update(pasteKey(p))
	run(t, m, cmd)

	if m.imgs.Len() != 1 {
		t.Fatalf("%d images attached, want 1", m.imgs.Len())
	}
	got := m.input.Value()
	if !strings.Contains(got, "[image 1]") {
		t.Errorf("prompt = %q", got)
	}
	if strings.Contains(got, p) {
		t.Errorf("the path was left in the prompt as well: %q", got)
	}
	// The status line names the file, so it is obvious which one was taken.
	if !strings.Contains(m.copied, "screenshot.png") {
		t.Errorf("status = %q", m.copied)
	}
}

// Pasted text is text, even when it mentions a file.
func TestPastingTextStillTypesIt(t *testing.T) {
	m := pasteModel(t, attach.Image{}, attach.ErrNoImage)

	_, cmd := m.Update(pasteKey("please fix internal/tui/model.go"))
	run(t, m, cmd)

	if m.imgs.Len() != 0 {
		t.Errorf("text became %d images", m.imgs.Len())
	}
	if got := m.input.Value(); got != "please fix internal/tui/model.go" {
		t.Errorf("prompt = %q", got)
	}
}

// A dropped file that cannot be used is worth a word — and the text still
// belongs in the box, because it is what was pasted.
func TestPastingAnUnreadableImageExplainsItself(t *testing.T) {
	p := filepath.Join(t.TempDir(), "broken.png")
	if err := os.WriteFile(p, []byte("not really a png"), 0o644); err != nil {
		t.Fatal(err)
	}
	m := pasteModel(t, attach.Image{}, attach.ErrNoImage)

	_, cmd := m.Update(pasteKey(p))
	run(t, m, cmd)

	if m.imgs.Len() != 0 {
		t.Fatalf("%d images attached from a broken file", m.imgs.Len())
	}
	if !strings.Contains(m.input.Value(), "broken.png") {
		t.Errorf("the pasted text was swallowed: %q", m.input.Value())
	}
	if !transcriptHas(m, "broken.png") {
		t.Error("nothing was said about the file that would not open")
	}
}

func transcriptHas(m *model, want string) bool {
	for _, e := range m.tr.entries {
		if strings.Contains(e.text, want) {
			return true
		}
	}
	return false
}

// ---- /image -------------------------------------------------------------

func TestImageCommandAttachesAPath(t *testing.T) {
	p := filepath.Join(t.TempDir(), "diagram.png")
	if err := os.WriteFile(p, testPNG(t), 0o644); err != nil {
		t.Fatal(err)
	}
	m := pasteModel(t, attach.Image{}, attach.ErrNoImage)

	run(t, m, m.attachPath(p))

	if m.imgs.Len() != 1 || !strings.Contains(m.input.Value(), "[image 1]") {
		t.Errorf("%d images, prompt %q", m.imgs.Len(), m.input.Value())
	}
}

func TestImageCommandOnAMissingFileComplains(t *testing.T) {
	m := pasteModel(t, attach.Image{}, attach.ErrNoImage)

	run(t, m, m.attachPath("/tmp/codey-no-such-picture.png"))

	if m.imgs.Len() != 0 {
		t.Fatal("a missing file became an image")
	}
	if !transcriptHas(m, "no image at") {
		t.Error("the missing file went unmentioned")
	}
}

// Bare /image is ctrl-V said out loud, so unlike the keystroke it reports an
// empty clipboard: it was asked for on purpose.
func TestBareImageCommandReportsAnEmptyClipboard(t *testing.T) {
	m := pasteModel(t, attach.Image{}, errors.New("no clipboard helper installed"))

	run(t, m, m.attachPath(""))

	if !transcriptHas(m, "clipboard helper") {
		t.Error("the missing helper went unmentioned")
	}
}

// ---- sending ------------------------------------------------------------

// The point of the whole feature: what goes up is a message with the picture
// in it, in the place the marker stood.
func TestSendingTurnsMarkersIntoImageBlocks(t *testing.T) {
	m := pasteModel(t, clipImage(t), nil)
	_, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlV})
	run(t, m, cmd)

	blocks := m.imgs.Blocks("what is wrong here " + m.input.Value())
	var img, text int
	for _, b := range blocks {
		switch b.Type {
		case llm.BlockImage:
			img++
			if b.Source == nil || b.Source.Data == "" {
				t.Error("an image block went up empty")
			}
		case llm.BlockText:
			text++
		}
	}
	if img != 1 || text != 1 {
		t.Fatalf("%d image and %d text blocks; want one of each", img, text)
	}
}

// The next prompt starts with an empty set, or its [image 1] would mean the
// last prompt's picture.
func TestRunTurnEmptiesTheSet(t *testing.T) {
	m := pasteModel(t, clipImage(t), nil)
	_, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlV})
	run(t, m, cmd)

	m.runTurn("have a look [image 1]") // the command itself is never run
	if m.imgs.Len() != 0 {
		t.Fatalf("%d images left over for the next prompt", m.imgs.Len())
	}
}

// Why ctrl+v is ours to take: the textarea binds it to a paste command of its
// own, but that command answers with a message type the bubbles package keeps
// to itself, and this model's Update has no branch that hands unknown
// messages back to the textarea. So the key was already doing nothing before
// it was intercepted — this test is what says so, and what will notice if a
// future Update starts forwarding.
func TestTheTextareasOwnPasteNeverArrives(t *testing.T) {
	m := newTestModel(t)
	before := m.input.Value()

	// Straight to the textarea, going around the interception.
	input, cmd := m.input.Update(tea.KeyMsg{Type: tea.KeyCtrlV})
	m.input = input
	if cmd == nil {
		t.Skip("the textarea no longer binds ctrl+v")
	}
	if msg := cmd(); msg != nil {
		m.Update(msg) // the runtime's route: back through the model
	}

	if m.input.Value() != before {
		t.Fatalf("the textarea's paste landed after all: %q", m.input.Value())
	}
}

// ctrl+v means paste. With no picture on the clipboard it has to fall back to
// the words on it, or taking the key over would have cost something.
func TestCtrlVWithoutAnImagePastesTheText(t *testing.T) {
	m := pasteModel(t, attach.Image{}, attach.ErrNoImage)
	m.textFn = func() (string, error) { return "some copied words", nil }
	m.input.SetValue("say ")

	_, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlV})
	run(t, m, cmd)

	if got := m.input.Value(); got != "say some copied words" {
		t.Errorf("prompt = %q", got)
	}
	if m.imgs.Len() != 0 {
		t.Errorf("text became %d images", m.imgs.Len())
	}
}

// A picture wins over the text form of the same clipboard: copying an image
// in a browser leaves both, and the picture is the one that was meant.
func TestAnImageBeatsTheClipboardsText(t *testing.T) {
	m := pasteModel(t, clipImage(t), nil)
	m.textFn = func() (string, error) { return "https://example.com/cat.png", nil }

	_, cmd := m.Update(tea.KeyMsg{Type: tea.KeyCtrlV})
	run(t, m, cmd)

	if got := m.input.Value(); !strings.Contains(got, "[image 1]") || strings.Contains(got, "http") {
		t.Errorf("prompt = %q", got)
	}
}

// The empty paste is a terminal that already did its own pasting and came up
// with nothing. Reading the clipboard's text again there could only duplicate
// what it just inserted, so that path asks for the picture alone.
func TestAnEmptyPasteDoesNotAlsoPasteText(t *testing.T) {
	m := pasteModel(t, attach.Image{}, attach.ErrNoImage)
	m.textFn = func() (string, error) {
		t.Error("the clipboard's text was read after an empty paste")
		return "", nil
	}

	_, cmd := m.Update(pasteKey(""))
	run(t, m, cmd)

	if got := m.input.Value(); got != "" {
		t.Errorf("prompt = %q, want it left alone", got)
	}
}

// typeRunes delivers text the way a terminal that does not bracket its pastes
// delivers it: as ordinary keystrokes, indistinguishable from typing.
func typeRunes(t *testing.T, m *model, s string) {
	t.Helper()
	for _, r := range s {
		m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{r}})
	}
}

// The bug this fixes: on a terminal without bracketed paste the path arrives
// as keystrokes, no paste is ever detected, and the prompt used to go up with
// a filename in it instead of a picture.
func TestAnUnbracketedPastedPathStillBecomesAnImage(t *testing.T) {
	p := filepath.Join(t.TempDir(), "clip.png")
	if err := os.WriteFile(p, testPNG(t), 0o644); err != nil {
		t.Fatal(err)
	}
	m := pasteModel(t, attach.Image{}, attach.ErrNoImage)

	typeRunes(t, m, `"`+p+`"`)
	if m.imgs.Len() != 0 {
		t.Fatal("typing attached something before enter was pressed")
	}

	m.Update(key("enter"))

	// runTurn has already spent the set, so the proof is what the transcript
	// shows: the marker, not the path.
	var said string
	for _, e := range m.tr.entries {
		if e.kind == entryUser {
			said = e.text
		}
	}
	if !strings.Contains(said, "[image 1]") {
		t.Errorf("the transcript shows %q, want a marker", said)
	}
	if strings.Contains(said, ".png") {
		t.Errorf("the path was left in the message: %q", said)
	}
}