Make SelectedLine/SelectedLineIdx work in staging/stagingSecondary views

While we try to keep the view's cursor position in sync with the context state's
selectedLineIdx (at least when pressing up or down), there are enough situations
where the two run out of sync; for example when initially opening the view, or
after staging a hunk, or when scrolling the view using the wheel. While it would
be possible to fix these situations to keep them always in sync, it doesn't seem
worth it, because the view's cursor position isn't really used for anything
else. So we rather special-case the SelectedLine/SelectedLineIdx functions of
ViewDriver to query the context state's selectedLineIdx directly if it is a
patch explorer context.

pkg/gui/gui_driver.go

@@ -49,6 +49,15 @@ func (self *GuiDriver) CurrentContext() types.Context {
 	return self.gui.c.CurrentContext()
 }
 
+func (self *GuiDriver) ContextForView(viewName string) types.Context {
+	context, ok := self.gui.contextForView(viewName)
+	if !ok {
+		return nil
+	}
+
+	return context
+}
+
 func (self *GuiDriver) Fail(message string) {
 	self.gui.g.Close()
 	// need to give the gui time to close

pkg/integration/components/test_test.go

@@ -30,6 +30,10 @@ func (self *fakeGuiDriver) CurrentContext() types.Context {
 	return nil
 }
 
+func (self *fakeGuiDriver) ContextForView(viewName string) types.Context {
+	return nil
+}
+
 func (self *fakeGuiDriver) Fail(message string) {
 	self.failureMessage = message
 }

pkg/integration/components/viewDriver.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/jesseduffield/gocui"
+	"github.com/jesseduffield/lazygit/pkg/gui/types"
 	"github.com/samber/lo"
 )
 
@@ -88,6 +89,9 @@ func (self *ViewDriver) Content(matcher *matcher) *ViewDriver {
 func (self *ViewDriver) SelectedLine(matcher *matcher) *ViewDriver {
 	self.t.matchString(matcher, fmt.Sprintf("%s: Unexpected selected line.", self.context),
 		func() string {
+			if idx, ok := self.selectedLineIdxInPatchExplorer(); ok {
+				return self.getView().BufferLines()[idx]
+			}
 			return self.getView().SelectedLine()
 		},
 	)
@@ -98,13 +102,25 @@ func (self *ViewDriver) SelectedLine(matcher *matcher) *ViewDriver {
 // asserts on the index of the selected line. 0 is the first index, representing the line at the top of the view.
 func (self *ViewDriver) SelectedLineIdx(expected int) *ViewDriver {
 	self.t.assertWithRetries(func() (bool, string) {
-		actual := self.getView().SelectedLineIdx()
+		actual, ok := self.selectedLineIdxInPatchExplorer()
+		if !ok {
+			actual = self.getView().SelectedLineIdx()
+		}
 		return expected == actual, fmt.Sprintf("%s: Expected selected line index to be %d, got %d", self.context, expected, actual)
 	})
 
 	return self
 }
 
+func (self *ViewDriver) selectedLineIdxInPatchExplorer() (int, bool) {
+	context := self.t.gui.ContextForView(self.getView().Name())
+	patchExplorerContext, ok := context.(types.IPatchExplorerContext)
+	if ok && patchExplorerContext.GetState() != nil {
+		return patchExplorerContext.GetState().GetSelectedLineIdx(), true
+	}
+	return 0, false
+}
+
 // focus the view (assumes the view is a side-view)
 func (self *ViewDriver) Focus() *ViewDriver {
 	viewName := self.getView().Name()

pkg/integration/tests/commit/unstaged.go

@@ -5,8 +5,6 @@ import (
 	. "github.com/jesseduffield/lazygit/pkg/integration/components"
 )
 
-// TODO: find out why we can't use .SelectedLine() on the staging/stagingSecondary views.
-
 var Unstaged = NewIntegrationTest(NewIntegrationTestArgs{
 	Description:  "Staging a couple files, going in the unstaged files menu, staging a line and committing",
 	ExtraCmdArgs: "",
@@ -30,11 +28,13 @@ var Unstaged = NewIntegrationTest(NewIntegrationTestArgs{
 			IsFocused().
 			Tap(func() {
 				t.Views().StagingSecondary().Content(DoesNotContain("+myfile content"))
+				t.Views().Staging().SelectedLine(Equals("+myfile content"))
 			}).
 			// stage the first line
 			PressPrimaryAction().
 			Tap(func() {
-				t.Views().Staging().Content(DoesNotContain("+myfile content"))
+				t.Views().Staging().Content(DoesNotContain("+myfile content")).
+					SelectedLine(Equals("+with a second line"))
 				t.Views().StagingSecondary().Content(Contains("+myfile content"))
 			}).
 			Press(keys.Files.CommitChanges)

pkg/integration/types/types.go

@@ -20,6 +20,7 @@ type GuiDriver interface {
 	PressKey(string)
 	Keys() config.KeybindingConfig
 	CurrentContext() types.Context
+	ContextForView(viewName string) types.Context
 	Fail(message string)
 	// These two log methods are for the sake of debugging while testing. There's no need to actually
 	// commit any logging.