Add option to (un)set upstream for a local branch

3 files changed, 111 insertions, 0 deletions

diff --git a/pkg/gui/controllers/helpers/helpers.go b/pkg/gui/controllers/helpers/helpers.go
index c45852e29..61fdedf46 100644
--- a/pkg/gui/controllers/helpers/helpers.go
+++ b/pkg/gui/controllers/helpers/helpers.go
@@ -12,6 +12,7 @@ type Helpers struct {
 	Host           *HostHelper
 	PatchBuilding  *PatchBuildingHelper
 	GPG            *GpgHelper
+	Upstream       *UpstreamHelper
 }
 
 func NewStubHelpers() *Helpers {
@@ -27,5 +28,6 @@ func NewStubHelpers() *Helpers {
 		Host:           &HostHelper{},
 		PatchBuilding:  &PatchBuildingHelper{},
 		GPG:            &GpgHelper{},
+		Upstream:       &UpstreamHelper{},
 	}
 }
diff --git a/pkg/gui/controllers/helpers/upstream_helper.go b/pkg/gui/controllers/helpers/upstream_helper.go
new file mode 100644
index 000000000..a0307a9d4
--- /dev/null
+++ b/pkg/gui/controllers/helpers/upstream_helper.go
@@ -0,0 +1,78 @@
+package helpers
+
+import (
+	"errors"
+	"strings"
+
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/jesseduffield/lazygit/pkg/gui/types"
+)
+
+type UpstreamHelper struct {
+	c     *types.HelperCommon
+	model *types.Model
+
+	getRemoteBranchesSuggestionsFunc func(string) func(string) []*types.Suggestion
+}
+
+type IUpstreamHelper interface {
+	ParseUpstream(string) (string, string, error)
+	PromptForUpstream(*models.Branch, func(string) error) error
+	GetSuggestedRemote() string
+}
+
+var _ IUpstreamHelper = &UpstreamHelper{}
+
+func NewUpstreamHelper(
+	c *types.HelperCommon,
+	model *types.Model,
+	getRemoteBranchesSuggestionsFunc func(string) func(string) []*types.Suggestion,
+) *UpstreamHelper {
+	return &UpstreamHelper{
+		c:                                c,
+		model:                            model,
+		getRemoteBranchesSuggestionsFunc: getRemoteBranchesSuggestionsFunc,
+	}
+}
+
+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
+}
+
+func (self *UpstreamHelper) PromptForUpstream(currentBranch *models.Branch, onConfirm func(string) error) error {
+	suggestedRemote := self.GetSuggestedRemote()
+
+	return self.c.Prompt(types.PromptOpts{
+		Title:               self.c.Tr.EnterUpstream,
+		InitialContent:      suggestedRemote + " " + currentBranch.Name,
+		FindSuggestionsFunc: self.getRemoteBranchesSuggestionsFunc(" "),
+		HandleConfirm:       onConfirm,
+	})
+}
+
+func (self *UpstreamHelper) GetSuggestedRemote() string {
+	return getSuggestedRemote(self.model.Remotes)
+}
+
+func getSuggestedRemote(remotes []*models.Remote) string {
+	if len(remotes) == 0 {
+		return "origin"
+	}
+
+	for _, remote := range remotes {
+		if remote.Name == "origin" {
+			return remote.Name
+		}
+	}
+
+	return remotes[0].Name
+}
diff --git a/pkg/gui/controllers/helpers/upstream_helper_test.go b/pkg/gui/controllers/helpers/upstream_helper_test.go
new file mode 100644
index 000000000..ac7a6a8bf
--- /dev/null
+++ b/pkg/gui/controllers/helpers/upstream_helper_test.go
@@ -0,0 +1,31 @@
+package helpers
+
+import (
+	"testing"
+
+	"github.com/jesseduffield/generics/slices"
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetSuggestedRemote(t *testing.T) {
+	cases := []struct {
+		remotes  []*models.Remote
+		expected string
+	}{
+		{mkRemoteList(), "origin"},
+		{mkRemoteList("upstream", "origin", "foo"), "origin"},
+		{mkRemoteList("upstream", "foo", "bar"), "upstream"},
+	}
+
+	for _, c := range cases {
+		result := getSuggestedRemote(c.remotes)
+		assert.EqualValues(t, c.expected, result)
+	}
+}
+
+func mkRemoteList(names ...string) []*models.Remote {
+	return slices.Map(names, func(name string) *models.Remote {
+		return &models.Remote{Name: name}
+	})
+}