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.