Merge pull request #1980 from ajhynes7/stash-untracked-changes

7 files changed, 124 insertions, 22 deletions

diff --git a/pkg/commands/git_commands/stash.go b/pkg/commands/git_commands/stash.go
index eddf676ff..572a87a73 100644
--- a/pkg/commands/git_commands/stash.go
+++ b/pkg/commands/git_commands/stash.go
@@ -123,6 +123,10 @@ func (self *StashCommands) SaveStagedChanges(message string) error {
 	return nil
 }
 
+func (self *StashCommands) StashIncludeUntrackedChanges(message string) error {
+	return self.cmd.New(fmt.Sprintf("git stash save %s --include-untracked", self.cmd.Quote(message))).Run()
+}
+
 func (self *StashCommands) Rename(index int, message string) error {
 	sha, err := self.Sha(index)
 	if err != nil {
diff --git a/pkg/gui/controllers/files_controller.go b/pkg/gui/controllers/files_controller.go
index 8f0b8333f..2a4434264 100644
--- a/pkg/gui/controllers/files_controller.go
+++ b/pkg/gui/controllers/files_controller.go
@@ -758,37 +758,53 @@ func (self *FilesController) createStashMenu() error {
 			{
 				Label: self.c.Tr.LcStashAllChanges,
 				OnPress: func() error {
-					return self.handleStashSave(self.git.Stash.Save, self.c.Tr.Actions.StashAllChanges, self.c.Tr.NoFilesToStash)
+					if !self.helpers.WorkingTree.IsWorkingTreeDirty() {
+						return self.c.ErrorMsg(self.c.Tr.NoFilesToStash)
+					}
+					return self.handleStashSave(self.git.Stash.Save, self.c.Tr.Actions.StashAllChanges)
 				},
 				Key: 'a',
 			},
 			{
 				Label: self.c.Tr.LcStashAllChangesKeepIndex,
 				OnPress: func() error {
+					if !self.helpers.WorkingTree.IsWorkingTreeDirty() {
+						return self.c.ErrorMsg(self.c.Tr.NoFilesToStash)
+					}
 					// if there are no staged files it behaves the same as Stash.Save
-					return self.handleStashSave(self.git.Stash.StashAndKeepIndex, self.c.Tr.Actions.StashAllChangesKeepIndex, self.c.Tr.NoFilesToStash)
+					return self.handleStashSave(self.git.Stash.StashAndKeepIndex, self.c.Tr.Actions.StashAllChangesKeepIndex)
 				},
 				Key: 'i',
 			},
 			{
+				Label: self.c.Tr.LcStashIncludeUntrackedChanges,
+				OnPress: func() error {
+					return self.handleStashSave(self.git.Stash.StashIncludeUntrackedChanges, self.c.Tr.Actions.StashIncludeUntrackedChanges)
+				},
+				Key: 'U',
+			},
+			{
 				Label: self.c.Tr.LcStashStagedChanges,
 				OnPress: func() error {
 					// there must be something in staging otherwise the current implementation mucks the stash up
 					if !self.helpers.WorkingTree.AnyStagedFiles() {
 						return self.c.ErrorMsg(self.c.Tr.NoTrackedStagedFilesStash)
 					}
-					return self.handleStashSave(self.git.Stash.SaveStagedChanges, self.c.Tr.Actions.StashStagedChanges, self.c.Tr.NoTrackedStagedFilesStash)
+					return self.handleStashSave(self.git.Stash.SaveStagedChanges, self.c.Tr.Actions.StashStagedChanges)
 				},
 				Key: 's',
 			},
 			{
 				Label: self.c.Tr.LcStashUnstagedChanges,
 				OnPress: func() error {
+					if !self.helpers.WorkingTree.IsWorkingTreeDirty() {
+						return self.c.ErrorMsg(self.c.Tr.NoFilesToStash)
+					}
 					if self.helpers.WorkingTree.AnyStagedFiles() {
-						return self.handleStashSave(self.git.Stash.StashUnstagedChanges, self.c.Tr.Actions.StashUnstagedChanges, self.c.Tr.NoFilesToStash)
+						return self.handleStashSave(self.git.Stash.StashUnstagedChanges, self.c.Tr.Actions.StashUnstagedChanges)
 					}
 					// ordinary stash
-					return self.handleStashSave(self.git.Stash.Save, self.c.Tr.Actions.StashUnstagedChanges, self.c.Tr.NoFilesToStash)
+					return self.handleStashSave(self.git.Stash.Save, self.c.Tr.Actions.StashUnstagedChanges)
 				},
 				Key: 'u',
 			},
@@ -797,7 +813,7 @@ func (self *FilesController) createStashMenu() error {
 }
 
 func (self *FilesController) stash() error {
-	return self.handleStashSave(self.git.Stash.Save, self.c.Tr.Actions.StashAllChanges, self.c.Tr.NoTrackedStagedFilesStash)
+	return self.handleStashSave(self.git.Stash.Save, self.c.Tr.Actions.StashAllChanges)
 }
 
 func (self *FilesController) createResetToUpstreamMenu() error {
@@ -825,11 +841,7 @@ func (self *FilesController) toggleTreeView() error {
 	return self.c.PostRefreshUpdate(self.context())
 }
 
-func (self *FilesController) handleStashSave(stashFunc func(message string) error, action string, errorMsg string) error {
-	if !self.helpers.WorkingTree.IsWorkingTreeDirty() {
-		return self.c.ErrorMsg(errorMsg)
-	}
-
+func (self *FilesController) handleStashSave(stashFunc func(message string) error, action string) error {
 	return self.c.Prompt(types.PromptOpts{
 		Title: self.c.Tr.StashChanges,
 		HandleConfirm: func(stashComment string) error {
diff --git a/pkg/i18n/english.go b/pkg/i18n/english.go
index 75e121545..e152b62a4 100644
--- a/pkg/i18n/english.go
+++ b/pkg/i18n/english.go
@@ -295,6 +295,7 @@ type TranslationSet struct {
 	LcStashStagedChanges                string
 	LcStashAllChangesKeepIndex          string
 	LcStashUnstagedChanges              string
+	LcStashIncludeUntrackedChanges      string
 	LcStashOptions                      string
 	NotARepository                      string
 	LcJump                              string
@@ -588,6 +589,7 @@ type Actions struct {
 	StashAllChangesKeepIndex          string
 	StashStagedChanges                string
 	StashUnstagedChanges              string
+	StashIncludeUntrackedChanges      string
 	GitFlowFinish                     string
 	GitFlowStart                      string
 	CopyToClipboard                   string
@@ -940,6 +942,7 @@ func EnglishTranslationSet() TranslationSet {
 		LcStashStagedChanges:                "stash staged changes",
 		LcStashAllChangesKeepIndex:          "stash all changes and keep index",
 		LcStashUnstagedChanges:              "stash unstaged changes",
+		LcStashIncludeUntrackedChanges:      "stash all changes including untracked files",
 		LcStashOptions:                      "Stash options",
 		NotARepository:                      "Error: must be run inside a git repository",
 		LcJump:                              "jump to panel",
@@ -1216,6 +1219,7 @@ func EnglishTranslationSet() TranslationSet {
 			StashAllChangesKeepIndex:          "Stash all changes and keep index",
 			StashStagedChanges:                "Stash staged changes",
 			StashUnstagedChanges:              "Stash unstaged changes",
+			StashIncludeUntrackedChanges:      "Stash all changes including untracked files",
 			GitFlowFinish:                     "Git flow finish",
 			GitFlowStart:                      "Git Flow start",
 			CopyToClipboard:                   "Copy to clipboard",
diff --git a/pkg/integration/components/assert.go b/pkg/integration/components/assert.go
index daa376c66..c911dbefa 100644
--- a/pkg/integration/components/assert.go
+++ b/pkg/integration/components/assert.go
@@ -84,6 +84,17 @@ func (self *Assert) CommitCount(expectedCount int) {
 	})
 }
 
+func (self *Assert) StashCount(expectedCount int) {
+	self.assertWithRetries(func() (bool, string) {
+		actualCount := len(self.gui.Model().StashEntries)
+
+		return actualCount == expectedCount, fmt.Sprintf(
+			"Expected %d stash entries, but got %d",
+			expectedCount, actualCount,
+		)
+	})
+}
+
 func (self *Assert) AtLeastOneCommit() {
 	self.assertWithRetries(func() (bool, string) {
 		actualCount := len(self.gui.Model().Commits)
diff --git a/pkg/integration/tests/stash/stash.go b/pkg/integration/tests/stash/stash.go
new file mode 100644
index 000000000..00d666168
--- /dev/null
+++ b/pkg/integration/tests/stash/stash.go
@@ -0,0 +1,34 @@
+package stash
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var Stash = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Stashing files",
+	ExtraCmdArgs: "",
+	Skip:         false,
+	SetupConfig:  func(config *config.AppConfig) {},
+	SetupRepo: func(shell *Shell) {
+		shell.EmptyCommit("initial commit")
+		shell.CreateFile("file", "content")
+		shell.GitAddAll()
+	},
+	Run: func(shell *Shell, input *Input, assert *Assert, keys config.KeybindingConfig) {
+		assert.StashCount(0)
+		assert.WorkingTreeFileCount(1)
+
+		input.PressKeys(keys.Files.ViewStashOptions)
+		assert.InMenu()
+
+		input.PressKeys("a")
+		assert.InPrompt()
+		assert.MatchCurrentViewTitle(Equals("Stash changes"))
+
+		input.Type("my stashed file")
+		input.Confirm()
+		assert.StashCount(1)
+		assert.WorkingTreeFileCount(0)
+	},
+})
diff --git a/pkg/integration/tests/stash/stash_including_untracked_files.go b/pkg/integration/tests/stash/stash_including_untracked_files.go
new file mode 100644
index 000000000..2f37f943c
--- /dev/null
+++ b/pkg/integration/tests/stash/stash_including_untracked_files.go
@@ -0,0 +1,35 @@
+package stash
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var StashIncludingUntrackedFiles = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Stashing all files including untracked ones",
+	ExtraCmdArgs: "",
+	Skip:         false,
+	SetupConfig:  func(config *config.AppConfig) {},
+	SetupRepo: func(shell *Shell) {
+		shell.EmptyCommit("initial commit")
+		shell.CreateFile("file_1", "content")
+		shell.CreateFile("file_2", "content")
+		shell.GitAdd("file_1")
+	},
+	Run: func(shell *Shell, input *Input, assert *Assert, keys config.KeybindingConfig) {
+		assert.StashCount(0)
+		assert.WorkingTreeFileCount(2)
+
+		input.PressKeys(keys.Files.ViewStashOptions)
+		assert.InMenu()
+
+		input.PressKeys("U")
+		assert.InPrompt()
+		assert.MatchCurrentViewTitle(Equals("Stash changes"))
+
+		input.Type("my stashed file")
+		input.Confirm()
+		assert.StashCount(1)
+		assert.WorkingTreeFileCount(0)
+	},
+})
diff --git a/pkg/integration/tests/tests.go b/pkg/integration/tests/tests.go
index 5c891a4e3..8d5f6adde 100644
--- a/pkg/integration/tests/tests.go
+++ b/pkg/integration/tests/tests.go
@@ -24,24 +24,26 @@ import (
 // be sure to add it to this list.
 
 var tests = []*components.IntegrationTest{
-	commit.Commit,
-	commit.NewBranch,
-	branch.Suggestions,
+	bisect.Basic,
+	bisect.FromOtherBranch,
 	branch.Delete,
 	branch.Rebase,
 	branch.RebaseAndDrop,
-	interactive_rebase.One,
-	interactive_rebase.AmendMerge,
-	custom_commands.Basic,
-	custom_commands.MultiplePrompts,
-	custom_commands.MenuFromCommand,
-	bisect.Basic,
-	bisect.FromOtherBranch,
+	branch.Suggestions,
 	cherry_pick.CherryPick,
 	cherry_pick.CherryPickConflicts,
+	commit.Commit,
+	commit.NewBranch,
+	custom_commands.Basic,
 	custom_commands.FormPrompts,
-	stash.Rename,
+	custom_commands.MenuFromCommand,
+	custom_commands.MultiplePrompts,
 	file.DirWithUntrackedFile,
+	interactive_rebase.AmendMerge,
+	interactive_rebase.One,
+	stash.Rename,
+	stash.Stash,
+	stash.StashIncludingUntrackedFiles,
 }
 
 func GetTests() []*components.IntegrationTest {