Vectrum
Sign In
465 lines · 14.8 KB · Go
package tui

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"

	tea "github.com/charmbracelet/bubbletea"

	"github.com/johannSo/codey/internal/agent"
	"github.com/johannSo/codey/internal/approval"
	"github.com/johannSo/codey/internal/llm"
	"github.com/johannSo/codey/internal/tools"
)

// pickerModel is a model wired the way main.go wires one — an agent, a session
// directory under the home root, and a factory that records which root the
// tools were last built against — with a list of trees already in hand, so the
// chooser can be driven without a git repository behind it.
func pickerModel(t *testing.T, trees ...worktreeInfo) (*model, *string) {
	t.Helper()
	root := t.TempDir()
	built := new(string)
	m := newModel(context.Background(), Options{
		Agent: agent.New(llm.New(llm.WithAPIKey("x")), tools.NewRegistry(nil),
			agent.NewState(), nil, agent.Options{}),
		Client:     llm.New(llm.WithAPIKey("x")),
		Policy:     approval.New(nil),
		Root:       root,
		HomeRoot:   root,
		SessionDir: filepath.Join(t.TempDir(), "chats"),
		SessionID:  "20260101T090000Z",
		NewTools: func(r string) *tools.Registry {
			*built = r
			return tools.NewRegistry(nil)
		},
	})
	m.Update(tea.WindowSizeMsg{Width: 100, Height: 30})
	m.worktrees = trees
	m.refilterWorktrees()
	return m, built
}

// tree makes a directory that can stand in for a worktree, since the chooser
// checks that what it is about to move into is still there.
func tree(t *testing.T, name string) worktreeInfo {
	t.Helper()
	path := filepath.Join(t.TempDir(), name)
	if err := os.MkdirAll(path, 0o755); err != nil {
		t.Fatal(err)
	}
	return worktreeInfo{Name: name, Path: path, Branch: name, Touched: time.Now()}
}

// openPicker puts the chooser up with the trees the model was given, without
// going near git.
func openPicker(m *model) {
	m.phase = phaseWorktrees
	m.worktreeFilter = ""
	m.refilterWorktrees()
	m.worktreeIdx = worktreeIndex(m.worktreeShown, m.opts.Root)
}

// The whole point of the subcommand: "resume" opens the chooser rather than
// being taken for the name of a tree to make.
func TestWorktreeResumeIsNotTakenForAName(t *testing.T) {
	m, built := pickerModel(t)
	m.opts.HomeRoot = t.TempDir() // not a repository, so listing fails loudly
	m.submit("/worktree resume")

	if m.worktreeing {
		t.Error("/worktree resume started making a tree called resume")
	}
	if *built != "" {
		t.Errorf("the tools were moved to %q by a failed listing", *built)
	}
	if out := m.tr.render(); !strings.Contains(out, "worktree") {
		t.Errorf("nothing was said about the failed listing:\n%s", out)
	}
}

func TestWorktreeResumeWithNoTreesSaysSo(t *testing.T) {
	if _, err := exec.LookPath("git"); err != nil {
		t.Skip("git is not installed")
	}
	isolateConfig(t)
	m, _ := pickerModel(t)
	repo := initRepo(t)
	m.opts.Root, m.opts.HomeRoot = repo, repo
	m.submit("/worktree resume")

	// git reports the checkout itself, which is not somewhere to go: a
	// one-row chooser offering where you already are is not a choice.
	if m.phase == phaseWorktrees {
		t.Error("a chooser was opened offering only the checkout itself")
	}
	if out := m.tr.render(); !strings.Contains(out, "no other worktree") {
		t.Errorf("a repository with no trees was not reported:\n%s", out)
	}
}

// Typing narrows the list, and what is left is what enter chooses from.
func TestWorktreePickerFiltersAsYouType(t *testing.T) {
	dark, parser, login := tree(t, "dark-mode"), tree(t, "parser-rewrite"), tree(t, "login-fix")
	m, _ := pickerModel(t, dark, parser, login)
	openPicker(m)

	typeIn(m, "par")
	if len(m.worktreeShown) != 1 || m.worktreeShown[0].Name != "parser-rewrite" {
		t.Fatalf("filtering by \"par\" left %v", treeNames(m.worktreeShown))
	}
	if view := m.worktreeView(); !strings.Contains(view, "parser-rewrite") || strings.Contains(view, "dark-mode") {
		t.Errorf("the view does not match the filter:\n%s", view)
	}

	// Backspace widens it again.
	pressKey(t, m, key("backspace"))
	pressKey(t, m, key("backspace"))
	pressKey(t, m, key("backspace"))
	if len(m.worktreeShown) != 3 {
		t.Errorf("after clearing the filter %v are shown, want all three", treeNames(m.worktreeShown))
	}

	// A filter nothing matches leaves an empty list that says so, and enter on
	// it does nothing rather than panicking on a row that is not there.
	typeIn(m, "zzz")
	if len(m.worktreeShown) != 0 {
		t.Fatalf("filtering by \"zzz\" left %v", treeNames(m.worktreeShown))
	}
	if view := m.worktreeView(); !strings.Contains(view, "nothing matches") {
		t.Errorf("an empty result was not explained:\n%s", view)
	}
	pressKey(t, m, key("enter"))
	if m.phase != phaseWorktrees {
		t.Error("enter on an empty list closed the chooser")
	}
}

// Choosing a tree moves the work into it and starts a fresh conversation: the
// tools are rebuilt there, the history is emptied, and the next turn is filed
// under a new session that records the tree.
func TestChoosingAWorktreeOpensANewChatInIt(t *testing.T) {
	dark := tree(t, "dark-mode")
	m, built := pickerModel(t, dark)

	m.opts.Agent.State().Replace([]llm.Message{user("the conversation being left")})
	m.turns = 1
	m.usage = llm.Usage{InputTokens: 500}
	before := m.opts.SessionID

	openPicker(m)
	pressKey(t, m, key("enter"))

	if m.phase != phaseIdle {
		t.Errorf("phase = %v, want the chooser closed", m.phase)
	}
	if m.opts.Root != dark.Path {
		t.Errorf("Root = %q, want the worktree %q", m.opts.Root, dark.Path)
	}
	if m.opts.Worktree != dark.Path {
		t.Errorf("Worktree = %q, want it recorded on the session", m.opts.Worktree)
	}
	if *built != dark.Path {
		t.Errorf("tools were built against %q, want the worktree", *built)
	}
	if got := m.opts.Agent.State().History(); len(got) != 0 {
		t.Errorf("the new chat started with %d messages, want none", len(got))
	}
	if m.turns != 0 || m.usage != (llm.Usage{}) {
		t.Errorf("the new chat started at %d turns and %+v, want nothing spent", m.turns, m.usage)
	}
	if m.opts.SessionID == before || m.opts.SessionID == "" {
		t.Errorf("session id = %q, want a new one", m.opts.SessionID)
	}

	// The conversation being left is on disk, so /resume can go back to it.
	left, err := agent.LoadSession(filepath.Join(m.opts.SessionDir, before+".json"))
	if err != nil {
		t.Fatalf("the conversation being left was not saved: %v", err)
	}
	if left.Title() != "the conversation being left" {
		t.Errorf("saved session title = %q", left.Title())
	}
	// Nothing is written for the chat that has said nothing yet.
	if _, err := os.Stat(filepath.Join(m.opts.SessionDir, m.opts.SessionID+".json")); err == nil {
		t.Error("an empty session was filed for a chat that has not started")
	}
	if out := m.tr.render(); !strings.Contains(out, "dark-mode") {
		t.Errorf("the move was not announced:\n%s", out)
	}
}

// Going back to the checkout Codey was started in is not moving into a tree,
// so the session must not record one — resuming it would otherwise send the
// work off into a "worktree" that is the main checkout.
func TestChoosingTheMainCheckoutClearsTheRecordedWorktree(t *testing.T) {
	m, _ := pickerModel(t)
	home := m.opts.HomeRoot
	dark := tree(t, "dark-mode")

	m.worktrees = []worktreeInfo{
		{Name: filepath.Base(home), Path: home, Branch: "main", Main: true},
		dark,
	}
	m.refilterWorktrees()

	m.openWorktreeChat(dark)
	if m.opts.Worktree != dark.Path {
		t.Fatalf("Worktree = %q, want the tree", m.opts.Worktree)
	}
	m.openWorktreeChat(m.worktrees[0])
	if m.opts.Root != home {
		t.Errorf("Root = %q, want the main checkout %q", m.opts.Root, home)
	}
	if m.opts.Worktree != "" {
		t.Errorf("Worktree = %q, want it cleared for the main checkout", m.opts.Worktree)
	}
}

// A tree removed behind Codey's back is reported rather than moved into.
func TestChoosingAWorktreeThatIsGoneIsRefused(t *testing.T) {
	gone := tree(t, "gone")
	if err := os.RemoveAll(gone.Path); err != nil {
		t.Fatal(err)
	}
	m, built := pickerModel(t, gone)
	before := m.opts.Root

	m.openWorktreeChat(gone)
	if m.opts.Root != before {
		t.Errorf("Root moved to %q, want it left alone", m.opts.Root)
	}
	if *built != "" {
		t.Errorf("tools were built against %q for a tree that is gone", *built)
	}
	if !strings.Contains(m.tr.render(), "is gone") {
		t.Errorf("a missing tree was not reported:\n%s", m.tr.render())
	}
}

// Moving mid-turn would pull the directory out from under the request in
// flight, so it is refused while anything is running.
func TestWorktreeResumeIsRefusedMidTurn(t *testing.T) {
	m, built := pickerModel(t, tree(t, "dark-mode"))
	m.phase = phaseRunning
	m.submit("/worktree resume")

	if m.phase != phaseRunning {
		t.Errorf("phase = %v, want the turn left alone", m.phase)
	}
	if *built != "" {
		t.Errorf("tools were moved to %q mid-turn", *built)
	}
	if !strings.Contains(m.tr.render(), "interrupt first") {
		t.Errorf("a mid-turn /worktree resume was not explained:\n%s", m.tr.render())
	}
}

// Esc leaves everything as it was.
func TestEscapingTheWorktreePickerChangesNothing(t *testing.T) {
	m, built := pickerModel(t, tree(t, "dark-mode"))
	before := m.opts.Root
	openPicker(m)
	typeIn(m, "da")
	pressKey(t, m, key("esc"))

	if m.phase != phaseIdle {
		t.Errorf("phase = %v, want the chooser closed", m.phase)
	}
	if m.opts.Root != before || *built != "" {
		t.Errorf("escaping moved the work to %q", m.opts.Root)
	}
	if m.worktreeFilter != "" || m.worktreeShown != nil {
		t.Error("the chooser kept its state after being closed")
	}
}

// A click on a row is a choice here too, and has to land on the tree it looks
// like it is on — the heading rows above the list are the easy thing to get
// wrong.
func TestClickingAWorktreeRowOpensIt(t *testing.T) {
	first, second := tree(t, "dark-mode"), tree(t, "parser-rewrite")
	m, _ := pickerModel(t, first, second)
	m.mouseOn = true
	openPicker(m)
	m.View()

	row, ok := m.rowOf(secWorktrees, worktreeHeadRows+1)
	if !ok {
		t.Fatal("the worktree chooser is not on screen anywhere")
	}
	m.Update(press(4, row))

	if m.opts.Root != second.Path {
		t.Errorf("clicking the second row opened %q, want %q", m.opts.Root, second.Path)
	}
	if m.phase != phaseIdle {
		t.Error("the chooser stayed open after a tree was clicked")
	}
}

// The wheel moves through this chooser rather than scrolling the transcript
// behind it, the way it does over every other open picker.
func TestTheWheelMovesThroughTheWorktreeChooser(t *testing.T) {
	m, _ := pickerModel(t, tree(t, "dark-mode"), tree(t, "parser-rewrite"))
	m.mouseOn = true
	openPicker(m)
	m.View()

	row, _ := m.rowOf(secWorktrees, worktreeHeadRows)
	m.Update(wheel(4, row, false))
	if m.worktreeIdx != 1 {
		t.Errorf("wheel down over the chooser moved to %d, want 1", m.worktreeIdx)
	}
	m.Update(wheel(4, row, true))
	if m.worktreeIdx != 0 {
		t.Errorf("wheel up over the chooser moved to %d, want 0", m.worktreeIdx)
	}
}

// git's porcelain output is the only description of the trees, so reading it
// has to survive a bare repository, a detached head, and a tree whose
// directory has been deleted without git being told.
func TestParseWorktreeList(t *testing.T) {
	home, dark := t.TempDir(), t.TempDir()
	out := fmt.Sprintf(`worktree %s
HEAD 1111111111111111111111111111111111111111
branch refs/heads/main

worktree %s
HEAD 2222222222222222222222222222222222222222
branch refs/heads/feature/dark-mode

worktree /nowhere/at/all
HEAD 3333333333333333333333333333333333333333
detached

worktree %s
bare
`, home, dark, t.TempDir())

	got := parseWorktreeList(out)
	if len(got) != 2 {
		t.Fatalf("parsed %v, want the main checkout and the one live tree", treeNames(got))
	}
	if !got[0].Main || got[0].Path != home || got[0].Branch != "main" {
		t.Errorf("first entry = %+v, want the main checkout", got[0])
	}
	if got[1].Main || got[1].Path != dark {
		t.Errorf("second entry = %+v, want the worktree", got[1])
	}
	if got[1].Branch != "feature/dark-mode" {
		t.Errorf("branch = %q, want the full name without refs/heads/", got[1].Branch)
	}
	if got[1].Name != filepath.Base(dark) {
		t.Errorf("name = %q, want the directory's", got[1].Name)
	}
}

// A bare repository is what git lists first when there is no main checkout,
// and it is nobody's checkout to go back to. Dropping it must not hand its
// place to the first real tree: a "main checkout" pinned to the top of the
// chooser would be offering a way home that does not exist.
func TestParseWorktreeListGivesABareRepoNoHeir(t *testing.T) {
	tree1, tree2 := t.TempDir(), t.TempDir()
	out := fmt.Sprintf(`worktree %s
bare

worktree %s
HEAD 1111111111111111111111111111111111111111
branch refs/heads/tree1

worktree %s
HEAD 2222222222222222222222222222222222222222
branch refs/heads/tree2
`, t.TempDir(), tree1, tree2)

	got := parseWorktreeList(out)
	if len(got) != 2 {
		t.Fatalf("parsed %v, want the two live trees", treeNames(got))
	}
	for _, w := range got {
		if w.Main {
			t.Errorf("%s calls itself the main checkout, but the repository is bare and has none", w.Name)
		}
	}
}

// A main checkout whose directory has been deleted without git being told is
// dropped, and it takes its position with it for the same reason.
func TestParseWorktreeListGivesADeletedCheckoutNoHeir(t *testing.T) {
	dark := t.TempDir()
	out := fmt.Sprintf(`worktree /nowhere/at/all
HEAD 1111111111111111111111111111111111111111
branch refs/heads/main

worktree %s
HEAD 2222222222222222222222222222222222222222
branch refs/heads/dark-mode
`, dark)

	got := parseWorktreeList(out)
	if len(got) != 1 {
		t.Fatalf("parsed %v, want only the tree that is still there", treeNames(got))
	}
	if got[0].Main {
		t.Errorf("%s inherited the deleted checkout's title", got[0].Name)
	}
}

// End to end against a real repository: a tree made by /worktree is the tree
// /worktree resume offers, and choosing it by name moves the work there.
func TestWorktreeResumeFindsATreeThatWasMade(t *testing.T) {
	if _, err := exec.LookPath("git"); err != nil {
		t.Skip("git is not installed")
	}
	isolateConfig(t)
	repo := initRepo(t)

	m, built := pickerModel(t)
	m.opts.Root, m.opts.HomeRoot = repo, repo
	m.opts.SessionDir = filepath.Join(t.TempDir(), "chats")

	path, err := runWorktreeScript(context.Background(), repo, "dark-mode")
	if err != nil {
		t.Fatal(err)
	}

	m.submit("/worktree resume")
	if m.phase != phaseWorktrees {
		t.Fatalf("phase = %v, want the chooser open:\n%s", m.phase, m.tr.render())
	}
	if view := m.worktreeView(); !strings.Contains(view, "dark-mode") {
		t.Errorf("the tree that was made is not in the chooser:\n%s", view)
	}

	pressKey(t, m, key("esc"))

	// By name, the chooser is skipped altogether.
	m.submit("/worktree resume dark-mode")
	if m.opts.Root != path {
		t.Errorf("Root = %q, want the worktree %q", m.opts.Root, path)
	}
	if *built != path {
		t.Errorf("tools were built against %q, want %q", *built, path)
	}

	m.submit("/worktree resume no-such-tree")
	if !strings.Contains(m.tr.render(), "no worktree called no-such-tree") {
		t.Errorf("an unknown name was not reported:\n%s", m.tr.render())
	}
}

func treeNames(trees []worktreeInfo) []string {
	out := make([]string, 0, len(trees))
	for _, w := range trees {
		out = append(out, w.Name)
	}
	return out
}