Standardize selection of upstream branches

pkg/gui/controllers.go

@@ -112,7 +112,7 @@ func (gui *Gui) resetHelpersAndControllers() {
 		MergeAndRebase:  rebaseHelper,
 		MergeConflicts:  mergeConflictsHelper,
 		CherryPick:      cherryPickHelper,
-		Upstream:        helpers.NewUpstreamHelper(helperCommon, suggestionsHelper.GetRemoteBranchesSuggestionsFunc),
+		Upstream:        helpers.NewUpstreamHelper(helperCommon, suggestionsHelper),
 		AmendHelper:     helpers.NewAmendHelper(helperCommon, gpgHelper),
 		FixupHelper:     helpers.NewFixupHelper(helperCommon),
 		Commits:         commitsHelper,

pkg/gui/controllers/branches_controller.go

@@ -273,26 +273,25 @@
 	setUpstreamItem := &types.MenuItem{
 		LabelColumns: []string{self.c.Tr.SetUpstream},
 		OnPress: func() error {
-			return self.c.Helpers().Upstream.PromptForUpstreamWithoutInitialContent(selectedBranch, func(upstream string) error {
-				upstreamRemote, upstreamBranch, err := self.c.Helpers().Upstream.ParseUpstream(upstream)
-				if err != nil {
-					return err
-				}
-
-				if err := self.c.Git().Branch.SetUpstream(upstreamRemote, upstreamBranch, selectedBranch.Name); err != nil {
-					return err
-				}
-				if err := self.c.Refresh(types.RefreshOptions{
-					Mode: types.SYNC,
-					Scope: []types.RefreshableView{
-						types.BRANCHES,
-						types.COMMITS,
-					},
-				}); err != nil {
-					return err
-				}
-				return nil
-			})
+			return self.c.Helpers().Upstream.PromptForUpstream(selectedBranch.Name,
+				func(remote string) string {
+					return fmt.Sprintf("Enter upstream branch on remote '%s'", remote)
+				},
+				func(upstream helpers.Upstream) error {
+					if err := self.c.Git().Branch.SetUpstream(upstream.Remote, upstream.Branch, selectedBranch.Name); err != nil {
+						return err
+					}
+					if err := self.c.Refresh(types.RefreshOptions{
+						Mode: types.SYNC,
+						Scope: []types.RefreshableView{
+							types.BRANCHES,
+							types.COMMITS,
+						},
+					}); err != nil {
+						return err
+					}
+					return nil
+				})
 		},
 		Key: 's',
 	}
@@ -756,21 +755,14 @@
 						return errors.New(self.c.Tr.PullRequestNoUpstream)
 					}
 
-					if len(self.c.Model().Remotes) == 1 {
-						toRemote := self.c.Model().Remotes[0].Name
-						self.c.Log.Debugf("PR will target the only existing remote '%s'", toRemote)
-						return self.promptForTargetBranchNameAndCreatePullRequest(branch, toRemote)
-					}
-
-					self.c.Prompt(types.PromptOpts{
-						Title:               self.c.Tr.SelectTargetRemote,
-						FindSuggestionsFunc: self.c.Helpers().Suggestions.GetRemoteSuggestionsFunc(),
-						HandleConfirm: func(toRemote string) error {
-							self.c.Log.Debugf("PR will target remote '%s'", toRemote)
-
-							return self.promptForTargetBranchNameAndCreatePullRequest(branch, toRemote)
+					self.c.Helpers().Upstream.PromptForUpstream("",
+						func(remote string) string {
+							return fmt.Sprintf("%s → %s/", branch.UpstreamBranch, remote)
 						},
-					})
+						func(upstream helpers.Upstream) error {
+							self.c.Log.Debugf("PR will target branch '%s' on remote '%s'", upstream.Branch, upstream.Remote)
+							return self.createPullRequest(branch.UpstreamBranch, upstream.Branch)
+						})
 
 					return nil
 				},
@@ -798,26 +790,6 @@
 	return self.c.Menu(types.CreateMenuOptions{Title: fmt.Sprint(self.c.Tr.CreatePullRequestOptions), Items: menuItems})
 }
 
-func (self *BranchesController) promptForTargetBranchNameAndCreatePullRequest(fromBranch *models.Branch, toRemote string) error {
-	remoteDoesNotExist := lo.NoneBy(self.c.Model().Remotes, func(remote *models.Remote) bool {
-		return remote.Name == toRemote
-	})
-	if remoteDoesNotExist {
-		return fmt.Errorf(self.c.Tr.NoValidRemoteName, toRemote)
-	}
-
-	self.c.Prompt(types.PromptOpts{
-		Title:               fmt.Sprintf("%s → %s/", fromBranch.UpstreamBranch, toRemote),
-		FindSuggestionsFunc: self.c.Helpers().Suggestions.GetRemoteBranchesForRemoteSuggestionsFunc(toRemote),
-		HandleConfirm: func(toBranch string) error {
-			self.c.Log.Debugf("PR will target branch '%s' on remote '%s'", toBranch, toRemote)
-			return self.createPullRequest(fromBranch.UpstreamBranch, toBranch)
-		},
-	})
-
-	return nil
-}
-
 func (self *BranchesController) createPullRequest(from string, to string) error {
 	url, err := self.c.Helpers().Host.GetPullRequestURL(from, to)
 	if err != nil {

pkg/gui/controllers/helpers/upstream_helper.go

@@ -1,75 +1,86 @@
 package helpers
 
 import (
-	"errors"
-	"strings"
+	"fmt"
 
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
+	"github.com/samber/lo"
 )
 
 type UpstreamHelper struct {
 	c *HelperCommon
 
-	getRemoteBranchesSuggestionsFunc func(string) func(string) []*types.Suggestion
+	suggestions *SuggestionsHelper
 }
 
 type IUpstreamHelper interface {
-	ParseUpstream(string) (string, string, error)
-	PromptForUpstreamWithInitialContent(*models.Branch, func(string) error) error
-	PromptForUpstreamWithoutInitialContent(*models.Branch, func(string) error) error
-	GetSuggestedRemote() string
+	PromptForUpstream(suggestedBranch string, branchSelectionTitle func(string) string, onConfirm func(Upstream) error) error
 }
 
 var _ IUpstreamHelper = &UpstreamHelper{}
 
 func NewUpstreamHelper(
 	c *HelperCommon,
-	getRemoteBranchesSuggestionsFunc func(string) func(string) []*types.Suggestion,
+	suggestions *SuggestionsHelper,
 ) *UpstreamHelper {
 	return &UpstreamHelper{
-		c:                                c,
-		getRemoteBranchesSuggestionsFunc: getRemoteBranchesSuggestionsFunc,
+		c:           c,
+		suggestions: suggestions,
 	}
 }
 
-func (self *UpstreamHelper) ParseUpstream(upstream string) (string, string, error) {
-	var upstreamBranch, upstreamRemote string
-	split := strings.Split(upstream, " ")
-	if len(split) != 2 {
-		return "", "", errors.New(self.c.Tr.InvalidUpstream)
-	}
-
-	upstreamRemote = split[0]
-	upstreamBranch = split[1]
-
-	return upstreamRemote, upstreamBranch, nil
+type Upstream struct {
+	Remote string
+	Branch string
 }
 
-func (self *UpstreamHelper) promptForUpstream(initialContent string, onConfirm func(string) error) error {
-	self.c.Prompt(types.PromptOpts{
-		Title:               self.c.Tr.EnterUpstream,
-		InitialContent:      initialContent,
-		FindSuggestionsFunc: self.getRemoteBranchesSuggestionsFunc(" "),
-		HandleConfirm:       onConfirm,
+func (self *UpstreamHelper) promptForUpstreamBranch(chosenRemote string, initialBranch string, branchSelectionTitle func(string) string, onConfirm func(Upstream) error) error {
+	remoteDoesNotExist := lo.NoneBy(self.c.Model().Remotes, func(remote *models.Remote) bool {
+		return remote.Name == chosenRemote
 	})
+	if remoteDoesNotExist {
+		return fmt.Errorf(self.c.Tr.NoValidRemoteName, chosenRemote)
+	}
 
+	self.c.Prompt(types.PromptOpts{
+		Title:               branchSelectionTitle(chosenRemote),
+		InitialContent:      initialBranch,
+		FindSuggestionsFunc: self.suggestions.GetRemoteBranchesForRemoteSuggestionsFunc(chosenRemote),
+		HandleConfirm: func(chosenBranch string) error {
+			self.c.Log.Debugf("User selected branch '%s' on remote '%s'", chosenRemote, chosenBranch)
+			return onConfirm(Upstream{chosenRemote, chosenBranch})
+		},
+	})
 	return nil
 }
 
-func (self *UpstreamHelper) PromptForUpstreamWithInitialContent(currentBranch *models.Branch, onConfirm func(string) error) error {
-	suggestedRemote := self.GetSuggestedRemote()
-	initialContent := suggestedRemote + " " + currentBranch.Name
-
-	return self.promptForUpstream(initialContent, onConfirm)
-}
-
-func (self *UpstreamHelper) PromptForUpstreamWithoutInitialContent(_ *models.Branch, onConfirm func(string) error) error {
-	return self.promptForUpstream("", onConfirm)
-}
+// Creates a series of two prompts that will gather an upstream remote and branch from the user.
+// The first prompt gathers the remote name, with a pre-filled default of origin, or the first remote in the list.
+// After selecting a remote, only branches present on that remote will be displayed.
+// If there is only one remote in the current repository, the remote prompt will be skipped.
+//
+// suggestedBranch pre-fills the second prompt, but can be an empty string if no suggestion would make sense. Often it is the current local branch name.
+// branchSelectionTitle allows customization of the second prompt to remind the user what their action entails.
+// onConfirm is called once the user has fully specified what their desired upstream is.
+func (self *UpstreamHelper) PromptForUpstream(suggestedBranch string, branchSelectionTitle func(string) string, onConfirm func(Upstream) error) error {
+	if len(self.c.Model().Remotes) == 1 {
+		remote := self.c.Model().Remotes[0].Name
+		self.c.Log.Debugf("Defaulting to only remote %s", remote)
+		return self.promptForUpstreamBranch(remote, suggestedBranch, branchSelectionTitle, onConfirm)
+	} else {
+		suggestedRemote := getSuggestedRemote(self.c.Model().Remotes)
+		self.c.Prompt(types.PromptOpts{
+			Title:               self.c.Tr.SelectTargetRemote,
+			InitialContent:      suggestedRemote,
+			FindSuggestionsFunc: self.suggestions.GetRemoteSuggestionsFunc(),
+			HandleConfirm: func(chosenRemote string) error {
+				return self.promptForUpstreamBranch(chosenRemote, suggestedBranch, branchSelectionTitle, onConfirm)
+			},
+		})
+	}
 
-func (self *UpstreamHelper) GetSuggestedRemote() string {
-	return getSuggestedRemote(self.c.Model().Remotes)
+	return nil
 }
 
 func getSuggestedRemote(remotes []*models.Remote) string {

pkg/gui/controllers/sync_controller.go

@@ -9,6 +9,7 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/commands/git_commands"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/context"
+	"github.com/jesseduffield/lazygit/pkg/gui/controllers/helpers"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 	"github.com/jesseduffield/lazygit/pkg/utils"
 )
@@ -99,18 +100,17 @@ func (self *SyncController) push(currentBranch *models.Branch) error {
 		if self.c.Git().Config.GetPushToCurrent() {
 			return self.pushAux(currentBranch, pushOpts{setUpstream: true})
 		} else {
-			return self.c.Helpers().Upstream.PromptForUpstreamWithInitialContent(currentBranch, func(upstream string) error {
-				upstreamRemote, upstreamBranch, err := self.c.Helpers().Upstream.ParseUpstream(upstream)
-				if err != nil {
-					return err
-				}
-
-				return self.pushAux(currentBranch, pushOpts{
-					setUpstream:    true,
-					upstreamRemote: upstreamRemote,
-					upstreamBranch: upstreamBranch,
+			return self.c.Helpers().Upstream.PromptForUpstream(currentBranch.Name,
+				func(remote string) string {
+					return fmt.Sprintf("Enter upstream branch on remote '%s'", remote)
+				},
+				func(upstream helpers.Upstream) error {
+					return self.pushAux(currentBranch, pushOpts{
+						setUpstream:    true,
+						upstreamRemote: upstream.Remote,
+						upstreamBranch: upstream.Branch,
+					})
 				})
-			})
 		}
 	}
 }
@@ -120,29 +120,28 @@ func (self *SyncController) pull(currentBranch *models.Branch) error {
 
 	// if we have no upstream branch we need to set that first
 	if !currentBranch.IsTrackingRemote() {
-		return self.c.Helpers().Upstream.PromptForUpstreamWithInitialContent(currentBranch, func(upstream string) error {
-			if err := self.setCurrentBranchUpstream(upstream); err != nil {
-				return err
-			}
+		return self.c.Helpers().Upstream.PromptForUpstream(currentBranch.Name,
+			func(remote string) string {
+				return fmt.Sprintf("Enter upstream branch on remote '%s'", remote)
+			},
+			func(upstream helpers.Upstream) error {
+				if err := self.setCurrentBranchUpstream(upstream); err != nil {
+					return err
+				}
 
-			return self.PullAux(currentBranch, PullFilesOptions{Action: action})
-		})
+				return self.PullAux(currentBranch, PullFilesOptions{Action: action})
+			})
 	}
 
 	return self.PullAux(currentBranch, PullFilesOptions{Action: action})
 }
 
-func (self *SyncController) setCurrentBranchUpstream(upstream string) error {
-	upstreamRemote, upstreamBranch, err := self.c.Helpers().Upstream.ParseUpstream(upstream)
-	if err != nil {
-		return err
-	}
-
-	if err := self.c.Git().Branch.SetCurrentBranchUpstream(upstreamRemote, upstreamBranch); err != nil {
+func (self *SyncController) setCurrentBranchUpstream(upstream helpers.Upstream) error {
+	if err := self.c.Git().Branch.SetCurrentBranchUpstream(upstream.Remote, upstream.Branch); err != nil {
 		if strings.Contains(err.Error(), "does not exist") {
 			return fmt.Errorf(
 				"upstream branch %s/%s not found.\nIf you expect it to exist, you should fetch (with 'f').\nOtherwise, you should push (with 'shift+P')",
-				upstreamRemote, upstreamBranch,
+				upstream.Remote, upstream.Branch,
 			)
 		}
 		return err

pkg/integration/tests/branch/open_pull_request_invalid_target_remote_name.go

@@ -42,6 +42,7 @@ var OpenPullRequestInvalidTargetRemoteName = NewIntegrationTest(NewIntegrationTe
 		t.ExpectPopup().
 			Prompt().
 			Title(Equals("Select target remote")).
+			Clear().
 			Type("non-existing-remote").
 			Confirm()
 

pkg/integration/tests/branch/set_upstream.go

@@ -28,9 +28,11 @@ var SetUpstream = NewIntegrationTest(NewIntegrationTestArgs{
 					Select(Contains(" Set upstream of selected branch")). // using leading space to disambiguate from the 'reset' option
 					Confirm()
 
+				// Origin is preselected on the users behalf because there is only one remote
+
 				t.ExpectPopup().Prompt().
-					Title(Equals("Enter upstream as '<remote> <branchname>'")).
-					SuggestionLines(Equals("origin master")).
+					Title(Equals("Enter upstream branch on remote 'origin'")).
+					SuggestionLines(Equals("master")).
 					ConfirmFirstSuggestion()
 			}).
 			Lines(

pkg/integration/tests/sync/pull_and_set_upstream.go

@@ -30,8 +30,8 @@ var PullAndSetUpstream = NewIntegrationTest(NewIntegrationTestArgs{
 		t.Views().Files().IsFocused().Press(keys.Universal.Pull)
 
 		t.ExpectPopup().Prompt().
-			Title(Equals("Enter upstream as '<remote> <branchname>'")).
-			SuggestionLines(Equals("origin master")).
+			Title(Equals("Enter upstream branch on remote 'origin'")).
+			SuggestionLines(Equals("master")).
 			ConfirmFirstSuggestion()
 
 		t.Views().Commits().

pkg/integration/tests/sync/push_and_set_upstream.go

@@ -26,8 +26,8 @@ var PushAndSetUpstream = NewIntegrationTest(NewIntegrationTestArgs{
 			Press(keys.Universal.Push)
 
 		t.ExpectPopup().Prompt().
-			Title(Equals("Enter upstream as '<remote> <branchname>'")).
-			SuggestionLines(Equals("origin master")).
+			Title(Equals("Enter upstream branch on remote 'origin'")).
+			SuggestionLines(Equals("master")).
 			ConfirmFirstSuggestion()
 
 		assertSuccessfullyPushed(t)