better handling of discarding files

4 files changed, 34 insertions, 2 deletions

diff --git a/pkg/commands/files.go b/pkg/commands/files.go
index 00b32e732..55a5f55d3 100644
--- a/pkg/commands/files.go
+++ b/pkg/commands/files.go
@@ -107,15 +107,34 @@ func (c *GitCommand) DiscardAllFileChanges(file *models.File) error {
 		return nil
 	}
 
-	// if the file isn't tracked, we assume you want to delete it
 	quotedFileName := c.OSCommand.Quote(file.Name)
+
+	if file.ShortStatus == "AA" {
+		if err := c.OSCommand.RunCommand("git checkout --ours --  %s", quotedFileName); err != nil {
+			return err
+		}
+		if err := c.OSCommand.RunCommand("git add %s", quotedFileName); err != nil {
+			return err
+		}
+		return nil
+	}
+
+	if file.ShortStatus == "DU" {
+		return c.OSCommand.RunCommand("git rm %s", quotedFileName)
+	}
+
+	// if the file isn't tracked, we assume you want to delete it
 	if file.HasStagedChanges || file.HasMergeConflicts {
 		if err := c.OSCommand.RunCommand("git reset -- %s", quotedFileName); err != nil {
 			return err
 		}
 	}
 
-	if !file.Tracked {
+	if file.ShortStatus == "DD" || file.ShortStatus == "AU" {
+		return nil
+	}
+
+	if file.Added {
 		return c.removeFile(file.Name)
 	}
 	return c.DiscardUnstagedFileChanges(file)
diff --git a/pkg/commands/git_test.go b/pkg/commands/git_test.go
index 1f290fbc0..9eea0c34c 100644
--- a/pkg/commands/git_test.go
+++ b/pkg/commands/git_test.go
@@ -344,6 +344,7 @@ func TestGitCommandGetStatusFiles(t *testing.T) {
 						HasStagedChanges:        true,
 						HasUnstagedChanges:      true,
 						Tracked:                 true,
+						Added:                   false,
 						Deleted:                 false,
 						HasMergeConflicts:       false,
 						HasInlineMergeConflicts: false,
@@ -356,6 +357,7 @@ func TestGitCommandGetStatusFiles(t *testing.T) {
 						HasStagedChanges:        true,
 						HasUnstagedChanges:      false,
 						Tracked:                 false,
+						Added:                   true,
 						Deleted:                 false,
 						HasMergeConflicts:       false,
 						HasInlineMergeConflicts: false,
@@ -368,6 +370,7 @@ func TestGitCommandGetStatusFiles(t *testing.T) {
 						HasStagedChanges:        true,
 						HasUnstagedChanges:      true,
 						Tracked:                 false,
+						Added:                   true,
 						Deleted:                 false,
 						HasMergeConflicts:       false,
 						HasInlineMergeConflicts: false,
@@ -380,6 +383,7 @@ func TestGitCommandGetStatusFiles(t *testing.T) {
 						HasStagedChanges:        false,
 						HasUnstagedChanges:      true,
 						Tracked:                 false,
+						Added:                   true,
 						Deleted:                 false,
 						HasMergeConflicts:       false,
 						HasInlineMergeConflicts: false,
@@ -392,6 +396,7 @@ func TestGitCommandGetStatusFiles(t *testing.T) {
 						HasStagedChanges:        false,
 						HasUnstagedChanges:      true,
 						Tracked:                 true,
+						Added:                   false,
 						Deleted:                 false,
 						HasMergeConflicts:       true,
 						HasInlineMergeConflicts: true,
@@ -1093,6 +1098,9 @@ func TestGitCommandUnstageFile(t *testing.T) {
 }
 
 // TestGitCommandDiscardAllFileChanges is a function.
+// these tests don't cover everything, in part because we already have an integration
+// test which does cover everything. I don't want to unnecessarily assert on the 'how'
+// when the 'what' is what matters
 func TestGitCommandDiscardAllFileChanges(t *testing.T) {
 	type scenario struct {
 		testName   string
@@ -1146,6 +1154,7 @@ func TestGitCommandDiscardAllFileChanges(t *testing.T) {
 			&models.File{
 				Name:    "test",
 				Tracked: false,
+				Added:   true,
 			},
 			func(string) error {
 				return fmt.Errorf("an error occurred when removing file")
@@ -1277,6 +1286,7 @@ func TestGitCommandDiscardAllFileChanges(t *testing.T) {
 			&models.File{
 				Name:             "test",
 				Tracked:          false,
+				Added:            true,
 				HasStagedChanges: true,
 			},
 			func(filename string) error {
@@ -1301,6 +1311,7 @@ func TestGitCommandDiscardAllFileChanges(t *testing.T) {
 			&models.File{
 				Name:             "test",
 				Tracked:          false,
+				Added:            true,
 				HasStagedChanges: false,
 			},
 			func(filename string) error {
diff --git a/pkg/commands/loading_files.go b/pkg/commands/loading_files.go
index 2bf77a865..f2a285f58 100644
--- a/pkg/commands/loading_files.go
+++ b/pkg/commands/loading_files.go
@@ -50,6 +50,7 @@ func (c *GitCommand) GetStatusFiles(opts GetStatusFileOptions) []*models.File {
 			HasUnstagedChanges:      unstagedChange != " ",
 			Tracked:                 !untracked,
 			Deleted:                 unstagedChange == "D" || stagedChange == "D",
+			Added:                   unstagedChange == "A" || untracked,
 			HasMergeConflicts:       hasMergeConflicts,
 			HasInlineMergeConflicts: hasInlineMergeConflicts,
 			Type:                    c.OSCommand.FileType(filename),
diff --git a/pkg/commands/models/file.go b/pkg/commands/models/file.go
index 482ba5cb5..4e20b7708 100644
--- a/pkg/commands/models/file.go
+++ b/pkg/commands/models/file.go
@@ -13,6 +13,7 @@ type File struct {
 	HasStagedChanges        bool
 	HasUnstagedChanges      bool
 	Tracked                 bool
+	Added                   bool
 	Deleted                 bool
 	HasMergeConflicts       bool
 	HasInlineMergeConflicts bool