From 8f8dd04ac7c27b796f1548397e6838de24a04d7d Mon Sep 17 00:00:00 2001
From: Jan De Dobbeleer <jan.de.dobbeleer@gmail.com>
Date: Mon, 4 Sep 2023 16:34:26 +0200
Subject: [PATCH] fix(git): correctly identify conflicted files

---
 src/segments/git.go      | 13 ++++++++++++-
 src/segments/git_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/src/segments/git.go b/src/segments/git.go
index 4dcfe4be..7450dc34 100644
--- a/src/segments/git.go
+++ b/src/segments/git.go
@@ -44,7 +44,7 @@ func (s *GitStatus) add(code string) {
 		s.Added++
 	case "?":
 		s.Untracked++
-	case "U":
+	case "U", "AA":
 		s.Unmerged++
 	case "M", "R", "C", "m":
 		s.Modified++
@@ -493,6 +493,17 @@ func (g *Git) setGitStatus() {
 		if len(status) <= 4 {
 			return
 		}
+
+		// map conflicts separately when in a merge or rebase
+		if g.Rebase || g.Merge {
+			conflict := "AA"
+			full := status[2:4]
+			if full == conflict {
+				g.Staging.add(conflict)
+				return
+			}
+		}
+
 		workingCode := status[3:4]
 		stagingCode := status[2:3]
 		g.Working.add(workingCode)
diff --git a/src/segments/git_test.go b/src/segments/git_test.go
index e7720067..e6f179dd 100644
--- a/src/segments/git_test.go
+++ b/src/segments/git_test.go
@@ -436,6 +436,8 @@ func TestSetGitStatus(t *testing.T) {
 		ExpectedUpstreamGone bool
 		ExpectedAhead        int
 		ExpectedBehind       int
+		Rebase               bool
+		Merge                bool
 	}{
 		{
 			Case: "all different options on working and staging, no remote",
@@ -535,6 +537,40 @@ func TestSetGitStatus(t *testing.T) {
 			ExpectedRef:          "branch-is-gone",
 			ExpectedUpstreamGone: true,
 		},
+		{
+			Case: "rebase with 2 merge conflicts",
+			Output: `
+			# branch.oid 1234567891011121314
+			# branch.head rework-git-status
+			# branch.upstream origin/rework-git-status
+			# branch.ab +0 -0
+			1 AA N...
+			1 AA N...
+			`,
+			ExpectedUpstream: "origin/rework-git-status",
+			ExpectedHash:     "1234567",
+			ExpectedRef:      "rework-git-status",
+			Rebase:           true,
+			ExpectedStaging:  &GitStatus{ScmStatus: ScmStatus{Unmerged: 2}},
+		},
+		{
+			Case: "merge with 4 merge conflicts",
+			Output: `
+			# branch.oid 1234567891011121314
+			# branch.head rework-git-status
+			# branch.upstream origin/rework-git-status
+			# branch.ab +0 -0
+			1 AA N...
+			1 AA N...
+			1 AA N...
+			1 AA N...
+			`,
+			ExpectedUpstream: "origin/rework-git-status",
+			ExpectedHash:     "1234567",
+			ExpectedRef:      "rework-git-status",
+			Merge:            true,
+			ExpectedStaging:  &GitStatus{ScmStatus: ScmStatus{Unmerged: 4}},
+		},
 	}
 	for _, tc := range cases {
 		env := new(mock.MockedEnvironment)
@@ -554,6 +590,8 @@ func TestSetGitStatus(t *testing.T) {
 		if tc.ExpectedStaging == nil {
 			tc.ExpectedStaging = &GitStatus{}
 		}
+		g.Rebase = tc.Rebase
+		g.Merge = tc.Merge
 		tc.ExpectedStaging.Formats = map[string]string{}
 		tc.ExpectedWorking.Formats = map[string]string{}
 		g.setGitStatus()