Add more worktree tests

12 files changed, 325 insertions, 30 deletions

diff --git a/docs/Custom_Command_Keybindings.md b/docs/Custom_Command_Keybindings.md
index 6b0a090ed..cca3985db 100644
--- a/docs/Custom_Command_Keybindings.md
+++ b/docs/Custom_Command_Keybindings.md
@@ -74,6 +74,7 @@ The permitted contexts are:
 | -------------- | -------------------------------------------------------------------------------------------------------- |
 | status         | The 'Status' tab                                                                                         |
 | files          | The 'Files' tab                                                                                          |
+| worktrees      | The 'Worktrees' tab                                                                                      |
 | localBranches  | The 'Local Branches' tab                                                                                 |
 | remotes        | The 'Remotes' tab                                                                                        |
 | remoteBranches | The context you get when pressing enter on a remote in the remotes tab                                   |
@@ -300,6 +301,7 @@ SelectedRemote
 SelectedTag
 SelectedStashEntry
 SelectedCommitFile
+SelectedWorktree
 CheckedOutBranch
 ```
 
diff --git a/pkg/commands/git_commands/bisect.go b/pkg/commands/git_commands/bisect.go
index 6deb32918..bd4b3ead2 100644
--- a/pkg/commands/git_commands/bisect.go
+++ b/pkg/commands/git_commands/bisect.go
@@ -19,12 +19,16 @@ func NewBisectCommands(gitCommon *GitCommon) *BisectCommands {
 // This command is pretty cheap to run so we're not storing the result anywhere.
 // But if it becomes problematic we can chang that.
 func (self *BisectCommands) GetInfo() *BisectInfo {
+	return self.GetInfoForGitDir(self.dotGitDir)
+}
+
+func (self *BisectCommands) GetInfoForGitDir(gitDir string) *BisectInfo {
 	var err error
 	info := &BisectInfo{started: false, log: self.Log, newTerm: "bad", oldTerm: "good"}
 	// we return nil if we're not in a git bisect session.
 	// we know we're in a session by the presence of a .git/BISECT_START file
 
-	bisectStartPath := filepath.Join(self.dotGitDir, "BISECT_START")
+	bisectStartPath := filepath.Join(gitDir, "BISECT_START")
 	exists, err := self.os.FileExists(bisectStartPath)
 	if err != nil {
 		self.Log.Infof("error getting git bisect info: %s", err.Error())
@@ -44,7 +48,7 @@ func (self *BisectCommands) GetInfo() *BisectInfo {
 	info.started = true
 	info.start = strings.TrimSpace(string(startContent))
 
-	termsContent, err := os.ReadFile(filepath.Join(self.dotGitDir, "BISECT_TERMS"))
+	termsContent, err := os.ReadFile(filepath.Join(gitDir, "BISECT_TERMS"))
 	if err != nil {
 		// old git versions won't have this file so we default to bad/good
 	} else {
@@ -53,7 +57,7 @@ func (self *BisectCommands) GetInfo() *BisectInfo {
 		info.oldTerm = splitContent[1]
 	}
 
-	bisectRefsDir := filepath.Join(self.dotGitDir, "refs", "bisect")
+	bisectRefsDir := filepath.Join(gitDir, "refs", "bisect")
 	files, err := os.ReadDir(bisectRefsDir)
 	if err != nil {
 		self.Log.Infof("error getting git bisect info: %s", err.Error())
@@ -85,7 +89,7 @@ func (self *BisectCommands) GetInfo() *BisectInfo {
 		info.statusMap[sha] = status
 	}
 
-	currentContent, err := os.ReadFile(filepath.Join(self.dotGitDir, "BISECT_EXPECTED_REV"))
+	currentContent, err := os.ReadFile(filepath.Join(gitDir, "BISECT_EXPECTED_REV"))
 	if err != nil {
 		self.Log.Infof("error getting git bisect info: %s", err.Error())
 		return info
diff --git a/pkg/commands/git_commands/worktree_loader.go b/pkg/commands/git_commands/worktree_loader.go
index b7e768ee0..6c73eaa13 100644
--- a/pkg/commands/git_commands/worktree_loader.go
+++ b/pkg/commands/git_commands/worktree_loader.go
@@ -48,10 +48,23 @@ func (self *WorktreeLoader) GetWorktrees() ([]*models.Worktree, error) {
 		}
 		if strings.HasPrefix(splitLine, "worktree ") {
 			path := strings.SplitN(splitLine, " ", 2)[1]
+			isMain := path == currentRepoPath
+
+			var gitDir string
+			if isMain {
+				gitDir = filepath.Join(path, ".git")
+			} else {
+				var ok bool
+				gitDir, ok = LinkedWorktreeGitPath(path)
+				if !ok {
+					self.Log.Warnf("Could not find git dir for worktree %s", path)
+				}
+			}
 
 			current = &models.Worktree{
 				IsMain: path == currentRepoPath,
 				Path:   path,
+				GitDir: gitDir,
 			}
 		} else if strings.HasPrefix(splitLine, "branch ") {
 			branch := strings.SplitN(splitLine, " ", 2)[1]
@@ -91,9 +104,21 @@ func (self *WorktreeLoader) GetWorktrees() ([]*models.Worktree, error) {
 			continue
 		}
 
+		// If we couldn't find the git directory, we can't find the branch name
+		if worktree.GitDir == "" {
+			continue
+		}
+
 		rebaseBranch, ok := rebaseBranch(worktree)
 		if ok {
 			worktree.Branch = rebaseBranch
+			continue
+		}
+
+		bisectBranch, ok := bisectBranch(worktree)
+		if ok {
+			worktree.Branch = bisectBranch
+			continue
 		}
 	}
 
@@ -101,29 +126,25 @@ func (self *WorktreeLoader) GetWorktrees() ([]*models.Worktree, error) {
 }
 
 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
+	for _, dir := range []string{"rebase-merge", "rebase-apply"} {
+		if bytesContent, err := os.ReadFile(filepath.Join(worktree.GitDir, dir, "head-name")); err == nil {
+			headName := strings.TrimSpace(string(bytesContent))
+			shortHeadName := strings.TrimPrefix(headName, "refs/heads/")
+			return shortHeadName, true
 		}
 	}
 
-	// now we look inside that git path for a file `rebase-merge/head-name`
-	// if it exists, we update the worktree to say that it has that for a head
-	headNameContents, err := os.ReadFile(filepath.Join(gitPath, "rebase-merge", "head-name"))
+	return "", false
+}
+
+func bisectBranch(worktree *models.Worktree) (string, bool) {
+	bisectStartPath := filepath.Join(worktree.GitDir, "BISECT_START")
+	startContent, err := os.ReadFile(bisectStartPath)
 	if err != nil {
 		return "", false
 	}
 
-	headName := strings.TrimSpace(string(headNameContents))
-	shortHeadName := strings.TrimPrefix(headName, "refs/heads/")
-
-	return shortHeadName, true
+	return strings.TrimSpace(string(startContent)), true
 }
 
 func LinkedWorktreeGitPath(worktreePath string) (string, bool) {
diff --git a/pkg/commands/models/worktree.go b/pkg/commands/models/worktree.go
index fb7dce62d..c14304233 100644
--- a/pkg/commands/models/worktree.go
+++ b/pkg/commands/models/worktree.go
@@ -4,9 +4,19 @@ package models
 type Worktree struct {
 	// if false, this is a linked worktree
 	IsMain bool
-	Path   string
+	// path to the directory of the worktree i.e. the directory that contains all the user's files
+	Path string
+	// path of the git directory for this worktree. The equivalent of the .git directory
+	// in the main worktree. For linked worktrees this would be <repo_path>/.git/worktrees/<name>
+	GitDir string
+	// If the worktree has a branch checked out, this field will be set to the branch name.
+	// A branch is considered 'checked out' if:
+	// * the worktree is directly on the branch
+	// * the worktree is mid-rebase on the branch
+	// * the worktree is mid-bisect on the branch
 	Branch string
-	// based on the path, but uniquified
+	// based on the path, but uniquified. Not the same name that git uses in the worktrees/ folder (no good reason for this,
+	// I just prefer my naming convention better)
 	NameField string
 }
 
diff --git a/pkg/gui/controllers/branches_controller.go b/pkg/gui/controllers/branches_controller.go
index 4ea8099b2..623514638 100644
--- a/pkg/gui/controllers/branches_controller.go
+++ b/pkg/gui/controllers/branches_controller.go
@@ -475,7 +475,10 @@ func (self *BranchesController) rename(branch *models.Branch) error {
 				}
 
 				// need to find where the branch is now so that we can re-select it. That means we need to refetch the branches synchronously and then find our branch
-				_ = self.c.Refresh(types.RefreshOptions{Mode: types.SYNC, Scope: []types.RefreshableView{types.BRANCHES}})
+				_ = self.c.Refresh(types.RefreshOptions{
+					Mode:  types.SYNC,
+					Scope: []types.RefreshableView{types.BRANCHES, types.WORKTREES},
+				})
 
 				// now that we've got our stuff again we need to find that branch and reselect it.
 				for i, newBranch := range self.c.Model().Branches {
diff --git a/pkg/gui/services/custom_commands/session_state_loader.go b/pkg/gui/services/custom_commands/session_state_loader.go
index 3566841b7..d5d34bfc9 100644
--- a/pkg/gui/services/custom_commands/session_state_loader.go
+++ b/pkg/gui/services/custom_commands/session_state_loader.go
@@ -33,6 +33,7 @@ type SessionState struct {
 	SelectedStashEntry     *models.StashEntry
 	SelectedCommitFile     *models.CommitFile
 	SelectedCommitFilePath string
+	SelectedWorktree       *models.Worktree
 	CheckedOutBranch       *models.Branch
 }
 
@@ -50,6 +51,7 @@ func (self *SessionStateLoader) call() *SessionState {
 		SelectedCommitFile:     self.c.Contexts().CommitFiles.GetSelectedFile(),
 		SelectedCommitFilePath: self.c.Contexts().CommitFiles.GetSelectedPath(),
 		SelectedSubCommit:      self.c.Contexts().SubCommits.GetSelected(),
+		SelectedWorktree:       self.c.Contexts().Worktrees.GetSelected(),
 		CheckedOutBranch:       self.refsHelper.GetCheckedOutRef(),
 	}
 }
diff --git a/pkg/integration/tests/test_list.go b/pkg/integration/tests/test_list.go
index e22ed7334..0e619f8b9 100644
--- a/pkg/integration/tests/test_list.go
+++ b/pkg/integration/tests/test_list.go
@@ -221,7 +221,11 @@ var tests = []*components.IntegrationTest{
 	undo.UndoCheckoutAndDrop,
 	undo.UndoDrop,
 	worktree.AddFromBranch,
+	worktree.AddFromBranchDetached,
+	worktree.AddFromCommit,
+	worktree.Bisect,
 	worktree.Crud,
+	worktree.CustomCommand,
 	worktree.DetachWorktreeFromBranch,
 	worktree.ForceRemoveWorktree,
 	worktree.Rebase,
diff --git a/pkg/integration/tests/worktree/add_from_branch_detached.go b/pkg/integration/tests/worktree/add_from_branch_detached.go
new file mode 100644
index 000000000..584de344e
--- /dev/null
+++ b/pkg/integration/tests/worktree/add_from_branch_detached.go
@@ -0,0 +1,46 @@
+package worktree
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var AddFromBranchDetached = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Add a detached 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 (detached)`)).
+					Confirm()
+
+				t.ExpectPopup().Prompt().
+					Title(Equals("New worktree path")).
+					Type("../linked-worktree").
+					Confirm()
+			}).
+			// confirm we're still focused on the branches view
+			IsFocused().
+			Lines(
+				Contains("(no branch)").IsSelected(),
+				Contains("mybranch (worktree)"),
+			)
+
+		t.Views().Status().
+			Content(Contains("repo(linked-worktree)"))
+	},
+})
diff --git a/pkg/integration/tests/worktree/add_from_commit.go b/pkg/integration/tests/worktree/add_from_commit.go
new file mode 100644
index 000000000..a171f74a3
--- /dev/null
+++ b/pkg/integration/tests/worktree/add_from_commit.go
@@ -0,0 +1,56 @@
+package worktree
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var AddFromCommit = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Add a worktree via the commits 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")
+		shell.EmptyCommit("commit two")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Commits().
+			Focus().
+			Lines(
+				Contains("commit two").IsSelected(),
+				Contains("initial commit"),
+			).
+			NavigateToLine(Contains("initial commit")).
+			Press(keys.Worktrees.ViewWorktreeOptions).
+			Tap(func() {
+				t.ExpectPopup().Menu().
+					Title(Equals("Worktree")).
+					Select(MatchesRegexp(`Create worktree from .*`).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()
+			}).
+			Lines(
+				Contains("initial commit"),
+			)
+
+		// Confirm we're now in the branches view
+		t.Views().Branches().
+			IsFocused().
+			Lines(
+				Contains("newbranch").IsSelected(),
+				Contains("mybranch (worktree)"),
+			)
+	},
+})
diff --git a/pkg/integration/tests/worktree/bisect.go b/pkg/integration/tests/worktree/bisect.go
new file mode 100644
index 000000000..143f8114f
--- /dev/null
+++ b/pkg/integration/tests/worktree/bisect.go
@@ -0,0 +1,88 @@
+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 bisect, 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.
+
+// not bothering to test the linked worktree here because it's the same logic as the rebase test
+
+var Bisect = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Verify that when you start a bisect in a linked 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").IsSelected(),
+				Contains("newbranch (worktree)"),
+			)
+
+		// start a bisect on the main worktree
+		t.Views().Commits().
+			Focus().
+			SelectedLine(Contains("commit 3")).
+			Press(keys.Commits.ViewBisectOptions).
+			Tap(func() {
+				t.ExpectPopup().Menu().
+					Title(Equals("Bisect")).
+					Select(MatchesRegexp(`Mark .* as bad`)).
+					Confirm()
+
+				t.Views().Information().Content(Contains("Bisecting"))
+			}).
+			NavigateToLine(Contains("initial commit")).
+			Press(keys.Commits.ViewBisectOptions).
+			Tap(func() {
+				t.ExpectPopup().Menu().
+					Title(Equals("Bisect")).
+					Select(MatchesRegexp(`Mark .* as good`)).
+					Confirm()
+			})
+
+		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("Bisecting"))
+			}).
+			Lines(
+				Contains("newbranch").IsSelected(),
+				Contains("mybranch (worktree)"),
+			)
+
+		// switch back to main worktree
+		t.Views().Branches().
+			Focus().
+			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()
+			})
+	},
+})
diff --git a/pkg/integration/tests/worktree/custom_command.go b/pkg/integration/tests/worktree/custom_command.go
new file mode 100644
index 000000000..2276e59be
--- /dev/null
+++ b/pkg/integration/tests/worktree/custom_command.go
@@ -0,0 +1,40 @@
+package worktree
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var CustomCommand = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Verify that custom commands work with worktrees by deleting a worktree via a custom command",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	SetupConfig: func(cfg *config.AppConfig) {
+		cfg.UserConfig.CustomCommands = []config.CustomCommand{
+			{
+				Key:     "d",
+				Context: "worktrees",
+				Command: "git worktree remove {{ .SelectedWorktree.Path | quote }}",
+			},
+		}
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.NewBranch("mybranch")
+		shell.CreateFileAndAdd("README.md", "hello world")
+		shell.Commit("initial commit")
+		shell.AddWorktree("mybranch", "../linked-worktree", "newbranch")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Worktrees().
+			Focus().
+			Lines(
+				Contains("repo (main)"),
+				Contains("linked-worktree"),
+			).
+			NavigateToLine(Contains("linked-worktree")).
+			Press("d").
+			Lines(
+				Contains("repo (main)"),
+			)
+	},
+})
diff --git a/pkg/integration/tests/worktree/rebase.go b/pkg/integration/tests/worktree/rebase.go
index e8ae59556..8b91702b5 100644
--- a/pkg/integration/tests/worktree/rebase.go
+++ b/pkg/integration/tests/worktree/rebase.go
@@ -10,8 +10,11 @@ import (
 // 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.
 
+// We need different logic for associated the branch depending on whether it's a main worktree or
+// linked worktree, so this test handles both.
+
 var Rebase = NewIntegrationTest(NewIntegrationTestArgs{
-	Description:  "Verify that when you start a rebase in a worktree, Lazygit still associates the worktree with the branch",
+	Description:  "Verify that when you start a rebase in a linked or main worktree, Lazygit still associates the worktree with the branch",
 	ExtraCmdArgs: []string{},
 	Skip:         false,
 	SetupConfig:  func(config *config.AppConfig) {},
@@ -27,10 +30,11 @@ var Rebase = NewIntegrationTest(NewIntegrationTestArgs{
 		t.Views().Branches().
 			Focus().
 			Lines(
-				Contains("mybranch"),
+				Contains("mybranch").IsSelected(),
 				Contains("newbranch (worktree)"),
 			)
 
+		// start a rebase on the main worktree
 		t.Views().Commits().
 			Focus().
 			NavigateToLine(Contains("commit 2")).
@@ -54,8 +58,19 @@ var Rebase = NewIntegrationTest(NewIntegrationTestArgs{
 			Lines(
 				Contains("newbranch").IsSelected(),
 				Contains("mybranch (worktree)"),
-			).
-			// switch back to main worktree
+			)
+
+		// start a rebase on the linked worktree
+		t.Views().Commits().
+			Focus().
+			NavigateToLine(Contains("commit 2")).
+			Press(keys.Universal.Edit)
+
+		t.Views().Information().Content(Contains("Rebasing"))
+
+		// switch back to main worktree
+		t.Views().Branches().
+			Focus().
 			NavigateToLine(Contains("mybranch")).
 			Press(keys.Universal.Select).
 			Tap(func() {
@@ -63,8 +78,12 @@ var Rebase = NewIntegrationTest(NewIntegrationTestArgs{
 					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"))
-			})
+			}).
+			Lines(
+				Contains("(no branch").IsSelected(),
+				Contains("mybranch"),
+				// even though the linked worktree is rebasing, we still associate it with the branch
+				Contains("newbranch (worktree)"),
+			)
 	},
 })