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.

5 files changed, 34 insertions, 4 deletions

diff --git a/pkg/gui/gui_driver.go b/pkg/gui/gui_driver.go
index cf1e3bbb4..239d1a43e 100644
--- a/pkg/gui/gui_driver.go
+++ b/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
diff --git a/pkg/integration/components/test_test.go b/pkg/integration/components/test_test.go
index 08823f525..c40fc3368 100644
--- a/pkg/integration/components/test_test.go
+++ b/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
 }
diff --git a/pkg/integration/components/viewDriver.go b/pkg/integration/components/viewDriver.go
index 04c58fc81..835d149db 100644
--- a/pkg/integration/components/viewDriver.go
+++ b/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()
diff --git a/pkg/integration/tests/commit/unstaged.go b/pkg/integration/tests/commit/unstaged.go
index c0e26b281..21e1e98d2 100644
--- a/pkg/integration/tests/commit/unstaged.go
+++ b/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)
diff --git a/pkg/integration/types/types.go b/pkg/integration/types/types.go
index 5326054d2..a26ac67af 100644
--- a/pkg/integration/types/types.go
+++ b/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.