support discarding unstaged changes

3 files changed, 51 insertions, 11 deletions

diff --git a/pkg/commands/git.go b/pkg/commands/git.go
index af39ab2d2..e674dc720 100644
--- a/pkg/commands/git.go
+++ b/pkg/commands/git.go
@@ -438,8 +438,8 @@ func (c *GitCommand) RebaseMode() (string, error) {
 	}
 }
 
-// RemoveFile directly
-func (c *GitCommand) RemoveFile(file *File) error {
+// DiscardAllFileChanges directly
+func (c *GitCommand) DiscardAllFileChanges(file *File) error {
 	// if the file isn't tracked, we assume you want to delete it
 	quotedFileName := c.OSCommand.Quote(file.Name)
 	if file.HasStagedChanges {
@@ -450,7 +450,12 @@ func (c *GitCommand) RemoveFile(file *File) error {
 	if !file.Tracked {
 		return c.removeFile(file.Name)
 	}
-	// if the file is tracked, we assume you want to just check it out
+	return c.DiscardUnstagedFileChanges(file)
+}
+
+// DiscardUnstagedFileChanges directly
+func (c *GitCommand) DiscardUnstagedFileChanges(file *File) error {
+	quotedFileName := c.OSCommand.Quote(file.Name)
 	return c.OSCommand.RunCommand(fmt.Sprintf("git checkout -- %s", quotedFileName))
 }
 
@@ -575,7 +580,7 @@ func (c *GitCommand) ApplyPatch(patch string) (string, error) {
 		return "", err
 	}
 
-	defer func() { _ = c.OSCommand.RemoveFile(filename) }()
+	defer func() { _ = c.OSCommand.Remove(filename) }()
 
 	return c.OSCommand.RunCommandWithOutput(fmt.Sprintf("git apply --cached %s", c.OSCommand.Quote(filename)))
 }
@@ -861,7 +866,7 @@ func (c *GitCommand) DiscardOldFileChanges(commits []*Commit, commitIndex int, f
 
 	// check if file exists in previous commit (this command returns an error if the file doesn't exist)
 	if err := c.OSCommand.RunCommand(fmt.Sprintf("git cat-file -e HEAD^:%s", fileName)); err != nil {
-		if err := c.OSCommand.RemoveFile(fileName); err != nil {
+		if err := c.OSCommand.Remove(fileName); err != nil {
 			return err
 		}
 		if err := c.StageFile(fileName); err != nil {
diff --git a/pkg/commands/git_test.go b/pkg/commands/git_test.go
index 8fbdc1642..614ad527b 100644
--- a/pkg/commands/git_test.go
+++ b/pkg/commands/git_test.go
@@ -1140,8 +1140,8 @@ func TestGitCommandIsInMergeState(t *testing.T) {
 	}
 }
 
-// TestGitCommandRemoveFile is a function.
-func TestGitCommandRemoveFile(t *testing.T) {
+// TestGitCommandDiscardAllFileChanges is a function.
+func TestGitCommandDiscardAllFileChanges(t *testing.T) {
 	type scenario struct {
 		testName   string
 		command    func() (func(string, ...string) *exec.Cmd, *[][]string)
@@ -1337,7 +1337,7 @@ func TestGitCommandRemoveFile(t *testing.T) {
 			gitCmd := NewDummyGitCommand()
 			gitCmd.OSCommand.command, cmdsCalled = s.command()
 			gitCmd.removeFile = s.removeFile
-			s.test(cmdsCalled, gitCmd.RemoveFile(s.file))
+			s.test(cmdsCalled, gitCmd.DiscardAllFileChanges(s.file))
 		})
 	}
 }
@@ -1932,3 +1932,38 @@ func TestGitCommandGetCommitFiles(t *testing.T) {
 		})
 	}
 }
+
+// TestGitCommandDiscardUnstagedChanges is a function.
+func TestGitCommandDiscardUnstagedChanges(t *testing.T) {
+	type scenario struct {
+		testName string
+		file     *File
+		command  func(string, ...string) *exec.Cmd
+		test     func(error)
+	}
+
+	scenarios := []scenario{
+		{
+			"valid case",
+			&File{Name: "test.txt"},
+			test.CreateMockCommand(t, []*test.CommandSwapper{
+				{
+					Expect:  `git checkout -- "test.txt"`,
+					Replace: "echo",
+				},
+			}),
+			func(err error) {
+				assert.NoError(t, err)
+			},
+		},
+	}
+
+	gitCmd := NewDummyGitCommand()
+
+	for _, s := range scenarios {
+		t.Run(s.testName, func(t *testing.T) {
+			gitCmd.OSCommand.command = s.command
+			s.test(gitCmd.DiscardUnstagedFileChanges(s.file))
+		})
+	}
+}
diff --git a/pkg/commands/os.go b/pkg/commands/os.go
index 9e0967cba..1fb8e8da6 100644
--- a/pkg/commands/os.go
+++ b/pkg/commands/os.go
@@ -255,9 +255,9 @@ func (c *OSCommand) CreateTempFile(filename, content string) (string, error) {
 	return tmpfile.Name(), nil
 }
 
-// RemoveFile removes a file at the specified path
-func (c *OSCommand) RemoveFile(filename string) error {
-	err := os.Remove(filename)
+// Remove removes a file or directory at the specified path
+func (c *OSCommand) Remove(filename string) error {
+	err := os.RemoveAll(filename)
 	return WrapError(err)
 }