Add worktree integration tests

3 files changed, 14 insertions, 9 deletions

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