Improve staging panel integration tests

11 files changed, 391 insertions, 99 deletions

diff --git a/pkg/integration/components/matcher.go b/pkg/integration/components/matcher.go
index 85a26bbde..ebec084a8 100644
--- a/pkg/integration/components/matcher.go
+++ b/pkg/integration/components/matcher.go
@@ -61,6 +61,11 @@ func (self *matcher) Contains(target string) *matcher {
 	return self.appendRule(matcherRule{
 		name: fmt.Sprintf("contains '%s'", target),
 		testFn: func(value string) (bool, string) {
+			// everything contains the empty string so we unconditionally return true here
+			if target == "" {
+				return true, ""
+			}
+
 			return strings.Contains(value, target), fmt.Sprintf("Expected '%s' to be found in '%s'", target, value)
 		},
 	})
diff --git a/pkg/integration/components/test_driver.go b/pkg/integration/components/test_driver.go
index 43618f23c..694a59dbf 100644
--- a/pkg/integration/components/test_driver.go
+++ b/pkg/integration/components/test_driver.go
@@ -7,7 +7,6 @@ import (
 
 	"github.com/atotto/clipboard"
 	"github.com/jesseduffield/lazygit/pkg/config"
-	"github.com/jesseduffield/lazygit/pkg/gui/types"
 	integrationTypes "github.com/jesseduffield/lazygit/pkg/integration/types"
 )
 
@@ -84,9 +83,7 @@ func (self *TestDriver) Shell() *Shell {
 // in the current page and failing that, jump to the top of the view and iterate through all of it,
 // looking for the item.
 func (self *TestDriver) navigateToListItem(matcher *matcher) {
-	self.inListContext()
-
-	currentContext := self.gui.CurrentContext().(types.IListContext)
+	currentContext := self.gui.CurrentContext()
 
 	view := currentContext.GetView()
 
@@ -133,14 +130,6 @@ func (self *TestDriver) navigateToListItem(matcher *matcher) {
 	}
 }
 
-func (self *TestDriver) inListContext() {
-	self.assertWithRetries(func() (bool, string) {
-		currentContext := self.gui.CurrentContext()
-		_, ok := currentContext.(types.IListContext)
-		return ok, fmt.Sprintf("Expected current context to be a list context, but got %s", currentContext.GetKey())
-	})
-}
-
 // for making assertions on lazygit views
 func (self *TestDriver) Views() *Views {
 	return &Views{t: self}
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..3f0b9726b 100644
--- a/pkg/integration/components/viewDriver.go
+++ b/pkg/integration/components/viewDriver.go
@@ -10,9 +10,10 @@ import (
 
 type ViewDriver struct {
 	// context is prepended to any error messages e.g. 'context: "current view"'
-	context string
-	getView func() *gocui.View
-	t       *TestDriver
+	context            string
+	getView            func() *gocui.View
+	t                  *TestDriver
+	getSelectedLinesFn func() ([]string, error)
 }
 
 // asserts that the view has the expected title
@@ -50,6 +51,74 @@ func (self *ViewDriver) Lines(matchers ...*matcher) *ViewDriver {
 	return self.assertLines(matchers...)
 }
 
+func (self *ViewDriver) getSelectedLines() ([]string, error) {
+	if self.getSelectedLinesFn == nil {
+		view := self.t.gui.View(self.getView().Name())
+
+		return []string{view.SelectedLine()}, nil
+	}
+
+	return self.getSelectedLinesFn()
+}
+
+func (self *ViewDriver) SelectedLines(matchers ...*matcher) *ViewDriver {
+	self.t.assertWithRetries(func() (bool, string) {
+		selectedLines, err := self.getSelectedLines()
+		if err != nil {
+			return false, err.Error()
+		}
+
+		selectedContent := strings.Join(selectedLines, "\n")
+		expectedContent := expectedContentFromMatchers(matchers)
+
+		if len(selectedLines) != len(matchers) {
+			return false, fmt.Sprintf("Expected the following to be selected:\n-----\n%s\n-----\nBut got:\n-----\n%s\n-----", expectedContent, selectedContent)
+		}
+
+		for i, line := range selectedLines {
+			ok, message := matchers[i].test(line)
+			if !ok {
+				return false, fmt.Sprintf("Error: %s. Expected the following to be selected:\n-----\n%s\n-----\nBut got:\n-----\n%s\n-----", message, expectedContent, selectedContent)
+			}
+		}
+
+		return true, ""
+	})
+
+	return self
+}
+
+func (self *ViewDriver) ContainsLines(matchers ...*matcher) *ViewDriver {
+	self.t.assertWithRetries(func() (bool, string) {
+		content := self.getView().Buffer()
+		lines := strings.Split(content, "\n")
+
+		for i := 0; i < len(lines)-len(matchers)+1; i++ {
+			matches := true
+			for j, matcher := range matchers {
+				ok, _ := matcher.test(lines[i+j])
+				if !ok {
+					matches = false
+					break
+				}
+			}
+			if matches {
+				return true, ""
+			}
+		}
+
+		expectedContent := expectedContentFromMatchers(matchers)
+
+		return false, fmt.Sprintf(
+			"Expected the following to be contained in the staging panel:\n-----\n%s\n-----\nBut got:\n-----\n%s\n-----",
+			expectedContent,
+			content,
+		)
+	})
+
+	return self
+}
+
 func (self *ViewDriver) assertLines(matchers ...*matcher) *ViewDriver {
 	view := self.getView()
 
@@ -86,9 +155,35 @@ func (self *ViewDriver) Content(matcher *matcher) *ViewDriver {
 
 // asserts on the selected line of the view
 func (self *ViewDriver) SelectedLine(matcher *matcher) *ViewDriver {
+	self.t.assertWithRetries(func() (bool, string) {
+		selectedLines, err := self.getSelectedLines()
+		if err != nil {
+			return false, err.Error()
+		}
+
+		if len(selectedLines) == 0 {
+			return false, "No line selected. Expected exactly one line to be selected"
+		} else if len(selectedLines) > 1 {
+			return false, fmt.Sprintf(
+				"Multiple lines selected. Expected only a single line to be selected. Selected lines:\n---\n%s\n---\n\nExpected line: %s",
+				strings.Join(selectedLines, "\n"),
+				matcher.name(),
+			)
+		}
+
+		value := selectedLines[0]
+		return matcher.context(fmt.Sprintf("%s: Unexpected selected line.", self.context)).test(value)
+	})
+
 	self.t.matchString(matcher, fmt.Sprintf("%s: Unexpected selected line.", self.context),
 		func() string {
-			return self.getView().SelectedLine()
+			selectedLines, err := self.getSelectedLines()
+			if err != nil {
+				self.t.gui.Fail(err.Error())
+				return "<failed to obtain selected line>"
+			}
+
+			return selectedLines[0]
 		},
 	)
 
@@ -253,3 +348,9 @@ func (self *ViewDriver) Tap(f func()) *ViewDriver {
 
 	return self
 }
+
+func expectedContentFromMatchers(matchers []*matcher) string {
+	return strings.Join(lo.Map(matchers, func(matcher *matcher, _ int) string {
+		return matcher.name()
+	}), "\n")
+}
diff --git a/pkg/integration/components/views.go b/pkg/integration/components/views.go
index ea5257602..1061d6262 100644
--- a/pkg/integration/components/views.go
+++ b/pkg/integration/components/views.go
@@ -2,8 +2,11 @@ package components
 
 import (
 	"fmt"
+	"strings"
 
+	"github.com/go-errors/errors"
 	"github.com/jesseduffield/gocui"
+	"github.com/jesseduffield/lazygit/pkg/gui/context"
 )
 
 type Views struct {
@@ -36,98 +39,129 @@ func (self *Views) Secondary() *ViewDriver {
 	}
 }
 
-func (self *Views) byName(viewName string) *ViewDriver {
+func (self *Views) regularView(viewName string) *ViewDriver {
+	return self.newStaticViewDriver(viewName, nil)
+}
+
+func (self *Views) patchExplorerViewByName(viewName string) *ViewDriver {
+	return self.newStaticViewDriver(viewName, func() ([]string, error) {
+		ctx := self.t.gui.ContextForView(viewName).(*context.PatchExplorerContext)
+		state := ctx.GetState()
+		if state == nil {
+			return nil, errors.New("Expected patch explorer to be activated")
+		}
+		selectedContent := state.PlainRenderSelected()
+		// the above method returns a string with a trailing newline so we need to remove that before splitting
+		selectedLines := strings.Split(strings.TrimSuffix(selectedContent, "\n"), "\n")
+		return selectedLines, nil
+	})
+}
+
+// 'static' because it'll always refer to the same view, as opposed to the 'main' view which could actually be
+// one of several views, or the 'current' view which depends on focus.
+func (self *Views) newStaticViewDriver(viewName string, getSelectedLinesFn func() ([]string, error)) *ViewDriver {
 	return &ViewDriver{
-		context: fmt.Sprintf("%s view", viewName),
-		getView: func() *gocui.View { return self.t.gui.View(viewName) },
-		t:       self.t,
+		context:            fmt.Sprintf("%s view", viewName),
+		getView:            func() *gocui.View { return self.t.gui.View(viewName) },
+		getSelectedLinesFn: getSelectedLinesFn,
+		t:                  self.t,
 	}
 }
 
 func (self *Views) Commits() *ViewDriver {
-	return self.byName("commits")
+	return self.regularView("commits")
 }
 
 func (self *Views) Files() *ViewDriver {
-	return self.byName("files")
+	return self.regularView("files")
 }
 
 func (self *Views) Status() *ViewDriver {
-	return self.byName("status")
+	return self.regularView("status")
 }
 
 func (self *Views) Submodules() *ViewDriver {
-	return self.byName("submodules")
+	return self.regularView("submodules")
 }
 
 func (self *Views) Information() *ViewDriver {
-	return self.byName("information")
+	return self.regularView("information")
 }
 
 func (self *Views) AppStatus() *ViewDriver {
-	return self.byName("appStatus")
+	return self.regularView("appStatus")
 }
 
 func (self *Views) Branches() *ViewDriver {
-	return self.byName("localBranches")
+	return self.regularView("localBranches")
 }
 
 func (self *Views) Remotes() *ViewDriver {
-	return self.byName("remotes")
+	return self.regularView("remotes")
 }
 
 func (self *Views) RemoteBranches() *ViewDriver {
-	return self.byName("remoteBranches")
+	return self.regularView("remoteBranches")
 }
 
 func (self *Views) Tags() *ViewDriver {
-	return self.byName("tags")
+	return self.regularView("tags")
 }
 
 func (self *Views) ReflogCommits() *ViewDriver {
-	return self.byName("reflogCommits")
+	return self.regularView("reflogCommits")
 }
 
 func (self *Views) SubCommits() *ViewDriver {
-	return self.byName("subCommits")
+	return self.regularView("subCommits")
 }
 
 func (self *Views) CommitFiles() *ViewDriver {
-	return self.byName("commitFiles")
+	return self.regularView("commitFiles")
 }
 
 func (self *Views) Stash() *ViewDriver {
-	return self.byName("stash")
+	return self.regularView("stash")
 }
 
 func (self *Views) Staging() *ViewDriver {
-	return self.byName("staging")
+	return self.patchExplorerViewByName("staging")
 }
 
 func (self *Views) StagingSecondary() *ViewDriver {
-	return self.byName("stagingSecondary")
+	return self.patchExplorerViewByName("stagingSecondary")
+}
+
+func (self *Views) PatchBuilding() *ViewDriver {
+	return self.patchExplorerViewByName("patchBuilding")
+}
+
+func (self *Views) PatchBuildingSecondary() *ViewDriver {
+	// this is not a patch explorer view because you can't actually focus it: it
+	// just renders content
+	return self.regularView("patchBuildingSecondary")
 }
 
 func (self *Views) Menu() *ViewDriver {
-	return self.byName("menu")
+	return self.regularView("menu")
 }
 
 func (self *Views) Confirmation() *ViewDriver {
-	return self.byName("confirmation")
+	return self.regularView("confirmation")
 }
 
 func (self *Views) CommitMessage() *ViewDriver {
-	return self.byName("commitMessage")
+	return self.regularView("commitMessage")
 }
 
 func (self *Views) Suggestions() *ViewDriver {
-	return self.byName("suggestions")
+	return self.regularView("suggestions")
 }
 
 func (self *Views) MergeConflicts() *ViewDriver {
-	return self.byName("mergeConflicts")
+	return self.regularView("mergeConflicts")
 }
 
 func (self *Views) Search() *ViewDriver {
-	return self.byName("search")
+	return self.regularView("search")
 }
diff --git a/pkg/integration/tests/staging/diff_context_change.go b/pkg/integration/tests/staging/diff_context_change.go
index 8ddf3d37f..a88b99848 100644
--- a/pkg/integration/tests/staging/diff_context_change.go
+++ b/pkg/integration/tests/staging/diff_context_change.go
@@ -50,31 +50,80 @@ var DiffContextChange = NewIntegrationTest(NewIntegrationTestArgs{
 
 		t.Views().Staging().
 			IsFocused().
-			Content(Contains("@@ -1,6 +1,6 @@").DoesNotContain(" 7a")).
-			SelectedLine(Contains("-3a")).
+			Press(keys.Main.ToggleSelectHunk).
+			SelectedLines(
+				Contains(`@@ -1,6 +1,6 @@`),
+				Contains(` 1a`),
+				Contains(` 2a`),
+				Contains(`-3a`),
+				Contains(`+3b`),
+				Contains(` 4a`),
+				Contains(` 5a`),
+				Contains(` 6a`),
+			).
 			Press(keys.Universal.IncreaseContextInDiffView).
-			// still on the same line
-			SelectedLine(Contains("-3a")).
-			// '7a' is now visible
-			Content(Contains("@@ -1,7 +1,7 @@").Contains(" 7a")).
+			SelectedLines(
+				Contains(`@@ -1,7 +1,7 @@`),
+				Contains(` 1a`),
+				Contains(` 2a`),
+				Contains(`-3a`),
+				Contains(`+3b`),
+				Contains(` 4a`),
+				Contains(` 5a`),
+				Contains(` 6a`),
+				Contains(` 7a`),
+			).
 			Press(keys.Universal.DecreaseContextInDiffView).
-			SelectedLine(Contains("-3a")).
-			Content(Contains("@@ -1,6 +1,6 @@").DoesNotContain(" 7a")).
+			SelectedLines(
+				Contains(`@@ -1,6 +1,6 @@`),
+				Contains(` 1a`),
+				Contains(` 2a`),
+				Contains(`-3a`),
+				Contains(`+3b`),
+				Contains(` 4a`),
+				Contains(` 5a`),
+				Contains(` 6a`),
+			).
 			Press(keys.Universal.DecreaseContextInDiffView).
-			SelectedLine(Contains("-3a")).
-			Content(Contains("@@ -1,5 +1,5 @@").DoesNotContain(" 6a")).
+			SelectedLines(
+				Contains(`@@ -1,5 +1,5 @@`),
+				Contains(` 1a`),
+				Contains(` 2a`),
+				Contains(`-3a`),
+				Contains(`+3b`),
+				Contains(` 4a`),
+				Contains(` 5a`),
+			).
 			Press(keys.Universal.DecreaseContextInDiffView).
-			// arguably we should still be on -3a, but at the moment the logic puts us on on +3b
-			SelectedLine(Contains("+3b")).
-			Content(Contains("@@ -2,3 +2,3 @@").DoesNotContain(" 5a")).
+			SelectedLines(
+				Contains(`@@ -2,3 +2,3 @@`),
+				Contains(` 2a`),
+				Contains(`-3a`),
+				Contains(`+3b`),
+				Contains(` 4a`),
+			).
 			PressPrimaryAction().
-			Content(DoesNotContain("+3b")).
 			Press(keys.Universal.TogglePanel)
 
 		t.Views().StagingSecondary().
 			IsFocused().
-			Content(Contains("@@ -3,2 +3,3 @@\n 3a\n+3b\n 4a")).
+			Press(keys.Main.ToggleSelectHunk).
+			SelectedLines(
+				Contains(`@@ -2,3 +2,3 @@`),
+				Contains(` 2a`),
+				Contains(`-3a`),
+				Contains(`+3b`),
+				Contains(` 4a`),
+			).
 			Press(keys.Universal.IncreaseContextInDiffView).
-			Content(Contains("@@ -2,4 +2,5 @@\n 2a\n 3a\n+3b\n 4a\n 5a"))
+			SelectedLines(
+				Contains(`@@ -1,5 +1,5 @@`),
+				Contains(` 1a`),
+				Contains(` 2a`),
+				Contains(`-3a`),
+				Contains(`+3b`),
+				Contains(` 4a`),
+				Contains(` 5a`),
+			)
 	},
 })
diff --git a/pkg/integration/tests/staging/discard_all_changes.go b/pkg/integration/tests/staging/discard_all_changes.go
index 7d27e7ec2..48cd0b62d 100644
--- a/pkg/integration/tests/staging/discard_all_changes.go
+++ b/pkg/integration/tests/staging/discard_all_changes.go
@@ -29,13 +29,13 @@ var DiscardAllChanges = NewIntegrationTest(NewIntegrationTestArgs{
 
 		t.Views().Staging().
 			IsFocused().
-			SelectedLine(Contains("+three")).
+			SelectedLines(Contains("+three")).
 			// discard the line
 			Press(keys.Universal.Remove).
 			Tap(func() {
 				t.Actions().ConfirmDiscardLines()
 			}).
-			SelectedLine(Contains("+four")).
+			SelectedLines(Contains("+four")).
 			// discard the other line
 			Press(keys.Universal.Remove).
 			Tap(func() {
@@ -49,6 +49,6 @@ var DiscardAllChanges = NewIntegrationTest(NewIntegrationTestArgs{
 			}).
 			// assert we are still in the staging panel, but now looking at the changes of the other file
 			IsFocused().
-			SelectedLine(Contains("+3"))
+			SelectedLines(Contains("+3"))
 	},
 })
diff --git a/pkg/integration/tests/staging/stage_hunks.go b/pkg/integration/tests/staging/stage_hunks.go
index b79819108..c2f85661a 100644
--- a/pkg/integration/tests/staging/stage_hunks.go
+++ b/pkg/integration/tests/staging/stage_hunks.go
@@ -50,42 +50,108 @@ var StageHunks = NewIntegrationTest(NewIntegrationTestArgs{
 
 		t.Views().Staging().
 			IsFocused().
-			SelectedLine(Contains("-3a")).
+			SelectedLines(
+				Contains("-3a"),
+			).
 			Press(keys.Universal.NextBlock).
-			SelectedLine(Contains("-13a")).
+			SelectedLines(
+				Contains("-13a"),
+			).
 			Press(keys.Main.ToggleSelectHunk).
+			SelectedLines(
+				Contains("@@ -10,6 +10,6 @@"),
+				Contains(" 10a"),
+				Contains(" 11a"),
+				Contains(" 12a"),
+				Contains("-13a"),
+				Contains("+13b"),
+				Contains(" 14a"),
+				Contains(" 15a"),
+				Contains(`\ No newline at end of file`),
+			).
 			// when in hunk mode, pressing up/down moves us up/down by a hunk
 			SelectPreviousItem().
-			SelectedLine(Contains("-3a")).
+			SelectedLines(
+				Contains(`@@ -1,6 +1,6 @@`),
+				Contains(` 1a`),
+				Contains(` 2a`),
+				Contains(`-3a`),
+				Contains(`+3b`),
+				Contains(` 4a`),
+				Contains(` 5a`),
+				Contains(` 6a`),
+			).
 			SelectNextItem().
-			SelectedLine(Contains("-13a")).
+			SelectedLines(
+				Contains("@@ -10,6 +10,6 @@"),
+				Contains(" 10a"),
+				Contains(" 11a"),
+				Contains(" 12a"),
+				Contains("-13a"),
+				Contains("+13b"),
+				Contains(" 14a"),
+				Contains(" 15a"),
+				Contains(`\ No newline at end of file`),
+			).
 			// stage the second hunk
 			PressPrimaryAction().
-			Content(Contains("-3a\n+3b")).
+			ContainsLines(
+				Contains("-3a"),
+				Contains("+3b"),
+			).
 			Tap(func() {
 				t.Views().StagingSecondary().
-					Content(Contains("-13a\n+13b"))
+					ContainsLines(
+						Contains("-13a"),
+						Contains("+13b"),
+					)
 			}).
 			Press(keys.Universal.TogglePanel)
 
 		t.Views().StagingSecondary().
 			IsFocused().
-			SelectedLine(Contains("-13a")).
 			// after toggling panel, we're back to only having selected a single line
+			SelectedLines(
+				Contains("-13a"),
+			).
 			PressPrimaryAction().
-			SelectedLine(Contains("+13b")).
+			SelectedLines(
+				Contains("+13b"),
+			).
 			PressPrimaryAction().
 			IsEmpty()
 
 		t.Views().Staging().
 			IsFocused().
-			SelectedLine(Contains("-3a")).
+			SelectedLines(
+				Contains("-3a"),
+			).
 			Press(keys.Main.ToggleSelectHunk).
+			SelectedLines(
+				Contains(`@@ -1,6 +1,6 @@`),
+				Contains(` 1a`),
+				Contains(` 2a`),
+				Contains(`-3a`),
+				Contains(`+3b`),
+				Contains(` 4a`),
+				Contains(` 5a`),
+				Contains(` 6a`),
+			).
 			Press(keys.Universal.Remove).
 			Tap(func() {
 				t.Actions().ConfirmDiscardLines()
 			}).
-			SelectedLine(Contains("-13a")).
-			Content(DoesNotContain("-3a").DoesNotContain("+3b"))
+			Content(DoesNotContain("-3a").DoesNotContain("+3b")).
+			SelectedLines(
+				Contains("@@ -10,6 +10,6 @@"),
+				Contains(" 10a"),
+				Contains(" 11a"),
+				Contains(" 12a"),
+				Contains("-13a"),
+				Contains("+13b"),
+				Contains(" 14a"),
+				Contains(" 15a"),
+				Contains(`\ No newline at end of file`),
+			)
 	},
 })
diff --git a/pkg/integration/tests/staging/stage_lines.go b/pkg/integration/tests/staging/stage_lines.go
index f1777fb36..f517440b8 100644
--- a/pkg/integration/tests/staging/stage_lines.go
+++ b/pkg/integration/tests/staging/stage_lines.go
@@ -26,16 +26,18 @@ var StageLines = NewIntegrationTest(NewIntegrationTestArgs{
 
 		t.Views().Staging().
 			IsFocused().
-			SelectedLine(Contains("+three")).
+			SelectedLines(Contains("+three")).
 			// stage 'three'
 			PressPrimaryAction().
 			// 'three' moves over to the staging secondary panel
 			Content(DoesNotContain("+three")).
 			Tap(func() {
 				t.Views().StagingSecondary().
-					Content(Contains("+three"))
+					ContainsLines(
+						Contains("+three"),
+					)
 			}).
-			SelectedLine(Contains("+four")).
+			SelectedLines(Contains("+four")).
 			// stage 'four'
 			PressPrimaryAction().
 			// nothing left in our staging panel
@@ -45,15 +47,20 @@ var StageLines = NewIntegrationTest(NewIntegrationTestArgs{
 		// do the same thing as above, moving the lines back to the staging panel
 		t.Views().StagingSecondary().
 			IsFocused().
-			Content(Contains("+three\n+four")).
-			SelectedLine(Contains("+three")).
+			ContainsLines(
+				Contains("+three"),
+				Contains("+four"),
+			).
+			SelectedLines(Contains("+three")).
 			PressPrimaryAction().
 			Content(DoesNotContain("+three")).
 			Tap(func() {
 				t.Views().Staging().
-					Content(Contains("+three"))
+					ContainsLines(
+						Contains("+three"),
+					)
 			}).
-			SelectedLine(Contains("+four")).
+			SelectedLines(Contains("+four")).
 			// pressing 'remove' has the same effect as pressing space when in the staging secondary panel
 			Press(keys.Universal.Remove).
 			IsEmpty()
@@ -61,8 +68,11 @@ var StageLines = NewIntegrationTest(NewIntegrationTestArgs{
 		// stage one line and then manually toggle to the staging secondary panel
 		t.Views().Staging().
 			IsFocused().
-			Content(Contains("+three\n+four")).
-			SelectedLine(Contains("+three")).
+			ContainsLines(
+				Contains("+three"),
+				Contains("+four"),
+			).
+			SelectedLines(Contains("+three")).
 			PressPrimaryAction().
 			Content(DoesNotContain("+three")).
 			Tap(func() {
@@ -77,7 +87,7 @@ var StageLines = NewIntegrationTest(NewIntegrationTestArgs{
 			Press(keys.Universal.TogglePanel)
 
 		t.Views().Staging().
-			SelectedLine(Contains("+four")).
+			SelectedLines(Contains("+four")).
 			// discard the line
 			Press(keys.Universal.Remove).
 			Tap(func() {
@@ -90,7 +100,9 @@ var StageLines = NewIntegrationTest(NewIntegrationTestArgs{
 
 		t.Views().StagingSecondary().
 			IsFocused().
-			Content(Contains("+three\n")).
+			ContainsLines(
+				Contains("+three"),
+			).
 			// return to file
 			PressEscape()
 
diff --git a/pkg/integration/tests/staging/stage_ranges.go b/pkg/integration/tests/staging/stage_ranges.go
index cd6d8021f..ea886ccda 100644
--- a/pkg/integration/tests/staging/stage_ranges.go
+++ b/pkg/integration/tests/staging/stage_ranges.go
@@ -26,50 +26,81 @@ var StageRanges = NewIntegrationTest(NewIntegrationTestArgs{
 
 		t.Views().Staging().
 			IsFocused().
-			SelectedLine(Contains("+three")).
+			SelectedLines(
+				Contains("+three"),
+			).
 			Press(keys.Main.ToggleDragSelect).
-			SelectNextItem().
-			SelectedLine(Contains("+four")).
-			SelectNextItem().
-			SelectedLine(Contains("+five")).
+			NavigateToListItem(Contains("+five")).
+			SelectedLines(
+				Contains("+three"),
+				Contains("+four"),
+				Contains("+five"),
+			).
 			// stage the three lines we've just selected
 			PressPrimaryAction().
-			Content(Contains(" five\n+six")).
+			SelectedLines(
+				Contains("+six"),
+			).
+			ContainsLines(
+				Contains(" five"),
+				Contains("+six"),
+			).
 			Tap(func() {
 				t.Views().StagingSecondary().
-					Content(Contains("+three\n+four\n+five"))
+					ContainsLines(
+						Contains("+three"),
+						Contains("+four"),
+						Contains("+five"),
+					)
 			}).
 			Press(keys.Universal.TogglePanel)
 
 		t.Views().StagingSecondary().
 			IsFocused().
-			SelectedLine(Contains("+three")).
+			SelectedLines(
+				Contains("+three"),
+			).
 			Press(keys.Main.ToggleDragSelect).
-			SelectNextItem().
-			SelectedLine(Contains("+four")).
-			SelectNextItem().
-			SelectedLine(Contains("+five")).
+			NavigateToListItem(Contains("+five")).
+			SelectedLines(
+				Contains("+three"),
+				Contains("+four"),
+				Contains("+five"),
+			).
 			// unstage the three selected lines
 			PressPrimaryAction().
 			// nothing left in our staging secondary panel
 			IsEmpty().
 			Tap(func() {
 				t.Views().Staging().
-					Content(Contains("+three\n+four\n+five\n+six"))
+					ContainsLines(
+						Contains("+three"),
+						Contains("+four"),
+						Contains("+five"),
+						Contains("+six"),
+					)
 			})
 
 		t.Views().Staging().
 			IsFocused().
 			// coincidentally we land at '+four' here. Maybe we should instead land
 			// at '+three'? given it's at the start of the hunk?
-			SelectedLine(Contains("+four")).
+			SelectedLines(
+				Contains("+four"),
+			).
 			Press(keys.Main.ToggleDragSelect).
 			SelectNextItem().
-			SelectedLine(Contains("+five")).
+			SelectedLines(
+				Contains("+four"),
+				Contains("+five"),
+			).
 			Press(keys.Universal.Remove).
 			Tap(func() {
 				t.Actions().ConfirmDiscardLines()
 			}).
-			Content(Contains("+three\n+six"))
+			ContainsLines(
+				Contains("+three"),
+				Contains("+six"),
+			)
 	},
 })
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.