Add worktree integration tests

12 files changed, 385 insertions, 18 deletions

diff --git a/pkg/commands/git_commands/worktree.go b/pkg/commands/git_commands/worktree.go
index b8c25c47b..65ab152a7 100644
--- a/pkg/commands/git_commands/worktree.go
+++ b/pkg/commands/git_commands/worktree.go
@@ -108,6 +108,8 @@ func CheckedOutByOtherWorktree(branch *models.Branch, worktrees []*models.Worktr
 	return !IsCurrentWorktree(worktree.Path)
 }
 
+// If in a non-bare repo, this returns the path of the main worktree
+// TODO: see if this works with a bare repo.
 func GetCurrentRepoPath() string {
 	pwd, err := os.Getwd()
 	if err != nil {
@@ -128,7 +130,7 @@ func GetCurrentRepoPath() string {
 	}
 
 	// either in a submodule, a worktree, or a bare repo
-	worktreeGitPath, ok := WorktreeGitPath(pwd)
+	worktreeGitPath, ok := LinkedWorktreeGitPath(pwd)
 	if !ok {
 		// fallback
 		return currentPath()
diff --git a/pkg/commands/git_commands/worktree_loader.go b/pkg/commands/git_commands/worktree_loader.go
index 0c863d546..b7e768ee0 100644
--- a/pkg/commands/git_commands/worktree_loader.go
+++ b/pkg/commands/git_commands/worktree_loader.go
@@ -28,6 +28,8 @@ func NewWorktreeLoader(
 }
 
 func (self *WorktreeLoader) GetWorktrees() ([]*models.Worktree, error) {
+	currentRepoPath := GetCurrentRepoPath()
+
 	cmdArgs := NewGitCmd("worktree").Arg("list", "--porcelain").ToArgv()
 	worktreesOutput, err := self.cmd.New(cmdArgs).DontLog().RunWithOutput()
 	if err != nil {
@@ -46,8 +48,9 @@ func (self *WorktreeLoader) GetWorktrees() ([]*models.Worktree, error) {
 		}
 		if strings.HasPrefix(splitLine, "worktree ") {
 			path := strings.SplitN(splitLine, " ", 2)[1]
+
 			current = &models.Worktree{
-				IsMain: len(worktrees) == 0,
+				IsMain: path == currentRepoPath,
 				Path:   path,
 			}
 		} else if strings.HasPrefix(splitLine, "branch ") {
@@ -88,7 +91,7 @@ func (self *WorktreeLoader) GetWorktrees() ([]*models.Worktree, error) {
 			continue
 		}
 
-		rebaseBranch, ok := rebaseBranch(worktree.Path)
+		rebaseBranch, ok := rebaseBranch(worktree)
 		if ok {
 			worktree.Branch = rebaseBranch
 		}
@@ -97,11 +100,17 @@ func (self *WorktreeLoader) GetWorktrees() ([]*models.Worktree, error) {
 	return worktrees, nil
 }
 
-func rebaseBranch(worktreePath string) (string, bool) {
-	// need to find the actual path of the worktree in the .git dir
-	gitPath, ok := WorktreeGitPath(worktreePath)
-	if !ok {
-		return "", false
+func rebaseBranch(worktree *models.Worktree) (string, bool) {
+	var gitPath string
+	if worktree.Main() {
+		gitPath = filepath.Join(worktree.Path, ".git")
+	} else {
+		// need to find the path of the linked worktree in the .git dir
+		var ok bool
+		gitPath, ok = LinkedWorktreeGitPath(worktree.Path)
+		if !ok {
+			return "", false
+		}
 	}
 
 	// now we look inside that git path for a file `rebase-merge/head-name`
@@ -117,7 +126,7 @@ func rebaseBranch(worktreePath string) (string, bool) {
 	return shortHeadName, true
 }
 
-func WorktreeGitPath(worktreePath string) (string, bool) {
+func LinkedWorktreeGitPath(worktreePath string) (string, bool) {
 	// first we get the path of the worktree, then we look at the contents of the `.git` file in that path
 	// then we look for the line that says `gitdir: /path/to/.git/worktrees/<worktree-name>`
 	// then we return that path
diff --git a/pkg/gui/context/worktrees_context.go b/pkg/gui/context/worktrees_context.go
index 7f15d67b1..17dd534a6 100644
--- a/pkg/gui/context/worktrees_context.go
+++ b/pkg/gui/context/worktrees_context.go
@@ -23,6 +23,7 @@ func NewWorktreesContext(c *ContextCommon) *WorktreesContext {
 
 	getDisplayStrings := func(startIdx int, length int) [][]string {
 		return presentation.GetWorktreeDisplayStrings(
+			c.Tr,
 			viewModel.GetFilteredList(),
 			c.Git().Worktree.IsCurrentWorktree,
 			c.Git().Worktree.IsWorktreePathMissing,
diff --git a/pkg/gui/controllers/helpers/worktree_helper.go b/pkg/gui/controllers/helpers/worktree_helper.go
index d75691525..aa4cea5e1 100644
--- a/pkg/gui/controllers/helpers/worktree_helper.go
+++ b/pkg/gui/controllers/helpers/worktree_helper.go
@@ -94,7 +94,7 @@ func (self *WorktreeHelper) NewWorktree() error {
 			HandleConfirm: func(base string) error {
 				// we assume that the base can be checked out
 				canCheckoutBase := true
-				return self.NewWorktreeCheckout(base, canCheckoutBase, detached)
+				return self.NewWorktreeCheckout(base, canCheckoutBase, detached, context.WORKTREES_CONTEXT_KEY)
 			},
 		})
 	}
@@ -120,7 +120,7 @@ func (self *WorktreeHelper) NewWorktree() error {
 	})
 }
 
-func (self *WorktreeHelper) NewWorktreeCheckout(base string, canCheckoutBase bool, detached bool) error {
+func (self *WorktreeHelper) NewWorktreeCheckout(base string, canCheckoutBase bool, detached bool, contextKey types.ContextKey) error {
 	opts := git_commands.NewWorktreeOpts{
 		Base:   base,
 		Detach: detached,
@@ -132,7 +132,7 @@ func (self *WorktreeHelper) NewWorktreeCheckout(base string, canCheckoutBase boo
 			if err := self.c.Git().Worktree.New(opts); err != nil {
 				return err
 			}
-			return self.Switch(opts.Path, context.LOCAL_BRANCHES_CONTEXT_KEY)
+			return self.Switch(opts.Path, contextKey)
 		})
 	}
 
@@ -251,13 +251,13 @@ func (self *WorktreeHelper) ViewBranchWorktreeOptions(branchName string, canChec
 			{
 				LabelColumns: []string{utils.ResolvePlaceholderString(self.c.Tr.CreateWorktreeFrom, placeholders)},
 				OnPress: func() error {
-					return self.NewWorktreeCheckout(branchName, canCheckoutBase, false)
+					return self.NewWorktreeCheckout(branchName, canCheckoutBase, false, context.LOCAL_BRANCHES_CONTEXT_KEY)
 				},
 			},
 			{
 				LabelColumns: []string{utils.ResolvePlaceholderString(self.c.Tr.CreateWorktreeFromDetached, placeholders)},
 				OnPress: func() error {
-					return self.NewWorktreeCheckout(branchName, canCheckoutBase, true)
+					return self.NewWorktreeCheckout(branchName, canCheckoutBase, true, context.LOCAL_BRANCHES_CONTEXT_KEY)
 				},
 			},
 		},
diff --git a/pkg/gui/presentation/worktrees.go b/pkg/gui/presentation/worktrees.go
index 4676a2847..6c8acf976 100644
--- a/pkg/gui/presentation/worktrees.go
+++ b/pkg/gui/presentation/worktrees.go
@@ -4,20 +4,22 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/presentation/icons"
 	"github.com/jesseduffield/lazygit/pkg/gui/style"
+	"github.com/jesseduffield/lazygit/pkg/i18n"
 	"github.com/jesseduffield/lazygit/pkg/theme"
 	"github.com/samber/lo"
 )
 
-func GetWorktreeDisplayStrings(worktrees []*models.Worktree, isCurrent func(string) bool, isMissing func(string) bool) [][]string {
+func GetWorktreeDisplayStrings(tr *i18n.TranslationSet, worktrees []*models.Worktree, isCurrent func(string) bool, isMissing func(string) bool) [][]string {
 	return lo.Map(worktrees, func(worktree *models.Worktree, _ int) []string {
 		return GetWorktreeDisplayString(
+			tr,
 			isCurrent(worktree.Path),
 			isMissing(worktree.Path),
 			worktree)
 	})
 }
 
-func GetWorktreeDisplayString(isCurrent bool, isPathMissing bool, worktree *models.Worktree) []string {
+func GetWorktreeDisplayString(tr *i18n.TranslationSet, isCurrent bool, isPathMissing bool, worktree *models.Worktree) []string {
 	textStyle := theme.DefaultTextColor
 
 	current := ""
@@ -41,8 +43,10 @@ func GetWorktreeDisplayString(isCurrent bool, isPathMissing bool, worktree *mode
 
 	name := worktree.Name()
 	if worktree.Main() {
-		// TODO: i18n
-		name += " (main worktree)"
+		name += " " + tr.MainWorktree
+	}
+	if isPathMissing && !icons.IsIconEnabled() {
+		name += " " + tr.MissingWorktree
 	}
 	res = append(res, textStyle.Sprint(name))
 	return res
diff --git a/pkg/integration/components/shell.go b/pkg/integration/components/shell.go
index decb748da..809cb1d5b 100644
--- a/pkg/integration/components/shell.go
+++ b/pkg/integration/components/shell.go
@@ -254,6 +254,13 @@ func (self *Shell) Init() *Shell {
 	return self
 }
 
+func (self *Shell) AddWorktree(base string, path string, newBranchName string) *Shell {
+	return self.RunCommand([]string{
+		"git", "worktree", "add", "-b",
+		newBranchName, path, base,
+	})
+}
+
 func (self *Shell) MakeExecutable(path string) *Shell {
 	// 0755 sets the executable permission for owner, and read/execute permissions for group and others
 	err := os.Chmod(filepath.Join(self.dir, path), 0o755)
diff --git a/pkg/integration/components/views.go b/pkg/integration/components/views.go
index 1a6e54b7e..eb4f585dc 100644
--- a/pkg/integration/components/views.go
+++ b/pkg/integration/components/views.go
@@ -129,6 +129,10 @@ func (self *Views) Files() *ViewDriver {
 	return self.regularView("files")
 }
 
+func (self *Views) Worktrees() *ViewDriver {
+	return self.regularView("worktrees")
+}
+
 func (self *Views) Status() *ViewDriver {
 	return self.regularView("status")
 }
diff --git a/pkg/integration/tests/test_list.go b/pkg/integration/tests/test_list.go
index 88d1883ac..d2e2848e7 100644
--- a/pkg/integration/tests/test_list.go
+++ b/pkg/integration/tests/test_list.go
@@ -26,6 +26,7 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/integration/tests/tag"
 	"github.com/jesseduffield/lazygit/pkg/integration/tests/ui"
 	"github.com/jesseduffield/lazygit/pkg/integration/tests/undo"
+	"github.com/jesseduffield/lazygit/pkg/integration/tests/worktree"
 )
 
 var tests = []*components.IntegrationTest{
@@ -219,4 +220,8 @@ var tests = []*components.IntegrationTest{
 	ui.SwitchTabFromMenu,
 	undo.UndoCheckoutAndDrop,
 	undo.UndoDrop,
+	worktree.AddFromBranch,
+	worktree.Crud,
+	worktree.Rebase,
+	worktree.WorktreeInRepo,
 }
diff --git a/pkg/integration/tests/worktree/add_from_branch.go b/pkg/integration/tests/worktree/add_from_branch.go
new file mode 100644
index 000000000..53636536d
--- /dev/null
+++ b/pkg/integration/tests/worktree/add_from_branch.go
@@ -0,0 +1,60 @@
+package worktree
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var AddFromBranch = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Add a worktree via the branches view, then switch back to the main worktree via the branches view",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig:  func(config *config.AppConfig) {},
+	SetupRepo: func(shell *Shell) {
+		shell.NewBranch("mybranch")
+		shell.CreateFileAndAdd("README.md", "hello world")
+		shell.Commit("initial commit")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Branches().
+			Focus().
+			Lines(
+				Contains("mybranch"),
+			).
+			Press(keys.Worktrees.ViewWorktreeOptions).
+			Tap(func() {
+				t.ExpectPopup().Menu().
+					Title(Equals("Worktree")).
+					Select(Contains(`Create worktree from mybranch`).DoesNotContain("detached")).
+					Confirm()
+
+				t.ExpectPopup().Prompt().
+					Title(Equals("New worktree path")).
+					Type("../linked-worktree").
+					Confirm()
+
+				t.ExpectPopup().Prompt().
+					Title(Equals("New branch name")).
+					Type("newbranch").
+					Confirm()
+			}).
+			// confirm we're still focused on the branches view
+			IsFocused().
+			Lines(
+				Contains("newbranch").IsSelected(),
+				Contains("mybranch (worktree)"),
+			).
+			NavigateToLine(Contains("mybranch")).
+			Press(keys.Universal.Select).
+			Tap(func() {
+				t.ExpectPopup().Confirmation().
+					Title(Equals("Switch to worktree")).
+					Content(Equals("This branch is checked out by worktree repo. Do you want to switch to that worktree?")).
+					Confirm()
+			}).
+			Lines(
+				Contains("mybranch").IsSelected(),
+				Contains("newbranch (worktree)"),
+			)
+	},
+})
diff --git a/pkg/integration/tests/worktree/crud.go b/pkg/integration/tests/worktree/crud.go
new file mode 100644
index 000000000..504d47b72
--- /dev/null
+++ b/pkg/integration/tests/worktree/crud.go
@@ -0,0 +1,120 @@
+package worktree
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var Crud = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "From the worktrees view, add a work tree, switch to it, switch back, and remove it",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig:  func(config *config.AppConfig) {},
+	SetupRepo: func(shell *Shell) {
+		shell.NewBranch("mybranch")
+		shell.CreateFileAndAdd("README.md", "hello world")
+		shell.Commit("initial commit")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Branches().
+			Lines(
+				Contains("mybranch"),
+			)
+
+		t.Views().Status().
+			Lines(
+				Contains("repo → mybranch"),
+			)
+
+		t.Views().Worktrees().
+			Focus().
+			Lines(
+				Contains("repo (main)"),
+			).
+			Press(keys.Universal.New).
+			Tap(func() {
+				t.ExpectPopup().Menu().
+					Title(Equals("Worktree")).
+					Select(Contains(`Create worktree from ref`).DoesNotContain(("detached"))).
+					Confirm()
+
+				t.ExpectPopup().Prompt().
+					Title(Equals("New worktree base ref")).
+					InitialText(Equals("mybranch")).
+					Confirm()
+
+				t.ExpectPopup().Prompt().
+					Title(Equals("New worktree path")).
+					Type("../linked-worktree").
+					Confirm()
+
+				t.ExpectPopup().Prompt().
+					Title(Equals("New branch name (leave blank to checkout mybranch)")).
+					Type("newbranch").
+					Confirm()
+			}).
+			Lines(
+				Contains("linked-worktree").IsSelected(),
+				Contains("repo (main)"),
+			).
+			// confirm we're still in the same view
+			IsFocused()
+
+		// status panel includes the worktree if it's a linked worktree
+		t.Views().Status().
+			Lines(
+				Contains("repo(linked-worktree) → newbranch"),
+			)
+
+		t.Views().Branches().
+			Lines(
+				Contains("newbranch"),
+				Contains("mybranch"),
+			)
+
+		t.Views().Worktrees().
+			// confirm we can't remove the current worktree
+			Press(keys.Universal.Remove).
+			Tap(func() {
+				t.ExpectPopup().Alert().
+					Title(Equals("Error")).
+					Content(Equals("You cannot remove the current worktree!")).
+					Confirm()
+			}).
+			// confirm we cannot remove the main worktree
+			NavigateToLine(Contains("repo (main)")).
+			Press(keys.Universal.Remove).
+			Tap(func() {
+				t.ExpectPopup().Alert().
+					Title(Equals("Error")).
+					Content(Equals("You cannot remove the main worktree!")).
+					Confirm()
+			}).
+			// switch back to main worktree
+			Press(keys.Universal.Select).
+			Lines(
+				Contains("repo (main)").IsSelected(),
+				Contains("linked-worktree"),
+			)
+
+		t.Views().Branches().
+			Lines(
+				Contains("mybranch"),
+				Contains("newbranch"),
+			)
+
+		t.Views().Worktrees().
+			// remove linked worktree
+			NavigateToLine(Contains("linked-worktree")).
+			Press(keys.Universal.Remove).
+			Tap(func() {
+				t.ExpectPopup().Confirmation().
+					Title(Equals("Remove worktree")).
+					Content(Contains("Are you sure you want to remove worktree 'linked-worktree'?")).
+					Confirm()
+			}).
+			Lines(
+				Contains("repo (main)").IsSelected(),
+			)
+	},
+})
diff --git a/pkg/integration/tests/worktree/rebase.go b/pkg/integration/tests/worktree/rebase.go
new file mode 100644
index 000000000..e8ae59556
--- /dev/null
+++ b/pkg/integration/tests/worktree/rebase.go
@@ -0,0 +1,70 @@
+package worktree
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+// This is important because `git worktree list` will show a worktree being in a detached head state (which is true)
+// when it's in the middle of a rebase, but it won't tell you about the branch it's on.
+// Even so, if you attempt to check out that branch from another worktree git won't let you, so we need to
+// keep track of the association ourselves.
+
+var Rebase = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Verify that when you start a rebase in a worktree, Lazygit still associates the worktree with the branch",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig:  func(config *config.AppConfig) {},
+	SetupRepo: func(shell *Shell) {
+		shell.NewBranch("mybranch")
+		shell.CreateFileAndAdd("README.md", "hello world")
+		shell.Commit("initial commit")
+		shell.EmptyCommit("commit 2")
+		shell.EmptyCommit("commit 3")
+		shell.AddWorktree("mybranch", "../linked-worktree", "newbranch")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Branches().
+			Focus().
+			Lines(
+				Contains("mybranch"),
+				Contains("newbranch (worktree)"),
+			)
+
+		t.Views().Commits().
+			Focus().
+			NavigateToLine(Contains("commit 2")).
+			Press(keys.Universal.Edit)
+
+		t.Views().Information().Content(Contains("Rebasing"))
+
+		t.Views().Branches().
+			Focus().
+			// switch to linked worktree
+			NavigateToLine(Contains("newbranch")).
+			Press(keys.Universal.Select).
+			Tap(func() {
+				t.ExpectPopup().Confirmation().
+					Title(Equals("Switch to worktree")).
+					Content(Equals("This branch is checked out by worktree linked-worktree. Do you want to switch to that worktree?")).
+					Confirm()
+
+				t.Views().Information().Content(DoesNotContain("Rebasing"))
+			}).
+			Lines(
+				Contains("newbranch").IsSelected(),
+				Contains("mybranch (worktree)"),
+			).
+			// switch back to main worktree
+			NavigateToLine(Contains("mybranch")).
+			Press(keys.Universal.Select).
+			Tap(func() {
+				t.ExpectPopup().Confirmation().
+					Title(Equals("Switch to worktree")).
+					Content(Equals("This branch is checked out by worktree repo. Do you want to switch to that worktree?")).
+					Confirm()
+
+				t.Views().Information().Content(Contains("Rebasing"))
+			})
+	},
+})
diff --git a/pkg/integration/tests/worktree/worktree_in_repo.go b/pkg/integration/tests/worktree/worktree_in_repo.go
new file mode 100644
index 000000000..743abddf5
--- /dev/null
+++ b/pkg/integration/tests/worktree/worktree_in_repo.go
@@ -0,0 +1,85 @@
+package worktree
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var WorktreeInRepo = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Add a worktree inside the repo, then remove the directory and confirm the worktree is removed",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig:  func(config *config.AppConfig) {},
+	SetupRepo: func(shell *Shell) {
+		shell.NewBranch("mybranch")
+		shell.CreateFileAndAdd("README.md", "hello world")
+		shell.Commit("initial commit")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Branches().
+			Lines(
+				Contains("mybranch"),
+			)
+
+		t.Views().Worktrees().
+			Focus().
+			Lines(
+				Contains("repo (main)"),
+			).
+			Press(keys.Universal.New).
+			Tap(func() {
+				t.ExpectPopup().Menu().
+					Title(Equals("Worktree")).
+					Select(Contains(`Create worktree from ref`).DoesNotContain(("detached"))).
+					Confirm()
+
+				t.ExpectPopup().Prompt().
+					Title(Equals("New worktree base ref")).
+					InitialText(Equals("mybranch")).
+					Confirm()
+
+				t.ExpectPopup().Prompt().
+					Title(Equals("New worktree path")).
+					Type("linked-worktree").
+					Confirm()
+
+				t.ExpectPopup().Prompt().
+					Title(Equals("New branch name (leave blank to checkout mybranch)")).
+					Type("newbranch").
+					Confirm()
+			}).
+			Lines(
+				Contains("linked-worktree").IsSelected(),
+				Contains("repo (main)"),
+			).
+			// switch back to main worktree
+			NavigateToLine(Contains("repo (main)")).
+			Press(keys.Universal.Select).
+			Lines(
+				Contains("repo (main)").IsSelected(),
+				Contains("linked-worktree"),
+			)
+
+		t.Views().Files().
+			Focus().
+			Lines(
+				Contains("linked-worktree"),
+			).
+			Press(keys.Universal.Remove).
+			Tap(func() {
+				t.ExpectPopup().Menu().
+					Title(Equals("linked-worktree")).
+					Select(Contains("Discard all changes")).
+					Confirm()
+			}).
+			IsEmpty()
+
+		// confirm worktree appears as missing
+		t.Views().Worktrees().
+			Focus().
+			Lines(
+				Contains("repo (main)").IsSelected(),
+				Contains("linked-worktree (missing)"),
+			)
+	},
+})