Rough draft for combining the coverage files

cmd/main.go

@@ -360,6 +360,15 @@ func goTestCmdArgs(opts *options, rerunOpts rerunOpts) []string {
 		result = append(result, rerunOpts.runFlag)
 	}
 
+	if rerunOpts.coverprofileFlag != "" {
+		// Replace the existing coverprofile arg with our new one in the re-run case.
+		coverprofileIndex, coverprofileIndexEnd := argIndex("coverprofile", args)
+		if coverprofileIndex >= 0 && coverprofileIndexEnd < len(args) {
+			args = append(args[:coverprofileIndex], args[coverprofileIndexEnd+1:]...)
+		}
+		result = append(result, rerunOpts.coverprofileFlag)
+	}
+
 	pkgArgIndex := findPkgArgPosition(args)
 	result = append(result, args[:pkgArgIndex]...)
 	result = append(result, cmdArgPackageList(opts, rerunOpts)...)

cmd/rerunfails.go

@@ -8,12 +8,14 @@ import (
 	"sort"
 	"strings"
 
+	"gotest.tools/gotestsum/internal/coverprofile"
 	"gotest.tools/gotestsum/testjson"
 )
 
 type rerunOpts struct {
-	runFlag string
-	pkg     string
+	runFlag          string
+	pkg              string
+	coverprofileFlag string
 }
 
 func (o rerunOpts) Args() []string {
@@ -24,13 +26,22 @@ func (o rerunOpts) Args() []string {
 	if o.pkg != "" {
 		result = append(result, o.pkg)
 	}
+	if o.coverprofileFlag != "" {
+		result = append(result, o.coverprofileFlag)
+	}
 	return result
 }
 
+func (o rerunOpts) withCoverprofile(coverprofile string) rerunOpts {
+	o.coverprofileFlag = "-coverprofile=" + coverprofile
+	return o
+}
+
 func newRerunOptsFromTestCase(tc testjson.TestCase) rerunOpts {
 	return rerunOpts{
-		runFlag: goTestRunFlagForTestCase(tc.Test),
-		pkg:     tc.Package,
+		runFlag:          goTestRunFlagForTestCase(tc.Test),
+		pkg:              tc.Package,
+		coverprofileFlag: "",
 	}
 }
 
@@ -51,22 +62,42 @@ func rerunFailsFilter(o *options) testCaseFilter {
 	return testjson.FilterFailedUnique
 }
 
+// Need to add the cov file context and etc...
 func rerunFailed(ctx context.Context, opts *options, scanConfig testjson.ScanConfig) error {
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 	tcFilter := rerunFailsFilter(opts)
 
+	// We need to take special care for the coverprofile file in the rerun
+	// failed case. If we pass the same `-coverprofile` flag to the `go test`
+	// command, it will overwrite the file. We need to combine the coverprofile
+	// files from the original run and the rerun.
+
+	isCoverprofile, mainProfile := coverprofile.ParseCoverProfile(opts.args)
+
 	rec := newFailureRecorderFromExecution(scanConfig.Execution)
 	for attempts := 0; rec.count() > 0 && attempts < opts.rerunFailsMaxAttempts; attempts++ {
 		testjson.PrintSummary(opts.stdout, scanConfig.Execution, testjson.SummarizeNone)
 		opts.stdout.Write([]byte("\n")) // nolint: errcheck
 
 		nextRec := newFailureRecorder(scanConfig.Handler)
 		for _, tc := range tcFilter(rec.failures) {
-			goTestProc, err := startGoTestFn(ctx, "", goTestCmdArgs(opts, newRerunOptsFromTestCase(tc)))
+			rerunOpts := newRerunOptsFromTestCase(tc)
+			var newCoverprofile string
+			if isCoverprofile {
+				// create a new coverprofile file
+				newCoverprofile = fmt.Sprintf("%s.%d", mainProfile, attempts)
+				rerunOpts = rerunOpts.withCoverprofile(newCoverprofile)
+			}
+			goTestProc, err := startGoTestFn(ctx, "", goTestCmdArgs(opts, rerunOpts))
 			if err != nil {
 				return err
 			}
+			if isCoverprofile {
+				if err := coverprofile.Combine(mainProfile, newCoverprofile); err != nil {
+					return err
+				}
+			}
 
 			cfg := testjson.ScanConfig{
 				RunID:     attempts + 1,

internal/coverprofile/covermerge.go

@@ -0,0 +1,95 @@
+package coverprofile
+
+// gocovmerge takes the results from multiple `go test -coverprofile` runs and
+// merges them into one profile
+
+// Taken from: https://github.com/wadey/gocovmerge @ b5bfa59ec0adc420475f97f89b58045c721d761c
+
+import (
+	"fmt"
+	"io"
+	"log"
+	"sort"
+
+	"golang.org/x/tools/cover"
+)
+
+func mergeProfiles(p *cover.Profile, merge *cover.Profile) {
+	if p.Mode != merge.Mode {
+		log.Fatalf("cannot merge profiles with different modes")
+	}
+	// Since the blocks are sorted, we can keep track of where the last block
+	// was inserted and only look at the blocks after that as targets for merge
+	startIndex := 0
+	for _, b := range merge.Blocks {
+		startIndex = mergeProfileBlock(p, b, startIndex)
+	}
+}
+
+func mergeProfileBlock(p *cover.Profile, pb cover.ProfileBlock, startIndex int) int {
+	sortFunc := func(i int) bool {
+		pi := p.Blocks[i+startIndex]
+		return pi.StartLine >= pb.StartLine && (pi.StartLine != pb.StartLine || pi.StartCol >= pb.StartCol)
+	}
+
+	i := 0
+	if !sortFunc(i) {
+		i = sort.Search(len(p.Blocks)-startIndex, sortFunc)
+	}
+	i += startIndex
+	if i < len(p.Blocks) && p.Blocks[i].StartLine == pb.StartLine && p.Blocks[i].StartCol == pb.StartCol {
+		if p.Blocks[i].EndLine != pb.EndLine || p.Blocks[i].EndCol != pb.EndCol {
+			log.Fatalf("OVERLAP MERGE: %v %v %v", p.FileName, p.Blocks[i], pb)
+		}
+		switch p.Mode {
+		case "set":
+			p.Blocks[i].Count |= pb.Count
+		case "count", "atomic":
+			p.Blocks[i].Count += pb.Count
+		default:
+			log.Fatalf("unsupported covermode: '%s'", p.Mode)
+		}
+	} else {
+		if i > 0 {
+			pa := p.Blocks[i-1]
+			if pa.EndLine >= pb.EndLine && (pa.EndLine != pb.EndLine || pa.EndCol > pb.EndCol) {
+				log.Fatalf("OVERLAP BEFORE: %v %v %v", p.FileName, pa, pb)
+			}
+		}
+		if i < len(p.Blocks)-1 {
+			pa := p.Blocks[i+1]
+			if pa.StartLine <= pb.StartLine && (pa.StartLine != pb.StartLine || pa.StartCol < pb.StartCol) {
+				log.Fatalf("OVERLAP AFTER: %v %v %v", p.FileName, pa, pb)
+			}
+		}
+		p.Blocks = append(p.Blocks, cover.ProfileBlock{})
+		copy(p.Blocks[i+1:], p.Blocks[i:])
+		p.Blocks[i] = pb
+	}
+	return i + 1
+}
+
+func addProfile(profiles []*cover.Profile, p *cover.Profile) []*cover.Profile {
+	i := sort.Search(len(profiles), func(i int) bool { return profiles[i].FileName >= p.FileName })
+	if i < len(profiles) && profiles[i].FileName == p.FileName {
+		mergeProfiles(profiles[i], p)
+	} else {
+		profiles = append(profiles, nil)
+		copy(profiles[i+1:], profiles[i:])
+		profiles[i] = p
+	}
+	return profiles
+}
+
+func dumpProfiles(profiles []*cover.Profile, out io.Writer) {
+	if len(profiles) == 0 {
+		return
+	}
+	fmt.Fprintf(out, "mode: %s\n", profiles[0].Mode)
+	for _, p := range profiles {
+		for _, b := range p.Blocks {
+			fmt.Fprintf(out, "%s:%d.%d,%d.%d %d %d\n", p.FileName, b.StartLine,
+				b.StartCol, b.EndLine, b.EndCol, b.NumStmt, b.Count)
+		}
+	}
+}

internal/coverprofile/coverprofile.go

@@ -0,0 +1,59 @@
+package coverprofile
+
+import (
+	"os"
+	"strings"
+
+	"golang.org/x/tools/cover"
+)
+
+// ParseCoverProfile parse the coverprofile file from the flag
+func ParseCoverProfile(args []string) (bool, string) {
+	for _, arg := range args {
+		if strings.HasPrefix(arg, "-coverprofile=") {
+			return true, strings.TrimPrefix(arg, "-coverprofile=")
+		}
+	}
+
+	return false, ""
+}
+
+func Combine(main string, other string) error {
+	var merged []*cover.Profile
+
+	mainProfiles, err := cover.ParseProfiles(main)
+	if err != nil {
+		return err
+	}
+
+	for _, p := range mainProfiles {
+		merged = addProfile(merged, p)
+	}
+
+	otherProfiles, err := cover.ParseProfiles(other)
+	if err != nil {
+		return err
+	}
+
+	for _, p := range otherProfiles {
+		merged = addProfile(merged, p)
+	}
+
+	// Create a tmp file to write the merged profiles to. Then use os.Rename to
+	// atomically move the file to the main profile to mimic the effect of
+	// atomic replacement of the file.
+	f, err := os.CreateTemp("", "gotestsum-coverprofile-rerun-*")
+	if err != nil {
+		return err
+	}
+	dumpProfiles(merged, f)
+	tempPath := f.Name()
+	if err := f.Close(); err != nil {
+		return err
+	}
+	if err := os.Rename(tempPath, main); err != nil {
+		return err
+	}
+
+	return nil
+}