From a979e497eb1df99e664120eb498aece9f9cddeba Mon Sep 17 00:00:00 2001
From: Jack Chistyakov <jchistyakov@meta.com>
Date: Mon, 20 Jan 2025 04:15:29 -0800
Subject: [PATCH] Make testmaingen more idiomatic

Summary: Better layout/organization. Needed for downstream work.

Reviewed By: inesusvet

Differential Revision: D68396871

fbshipit-source-id: 6e6976d724d985c144dbb2d73b32e636aa70182f
---
 go_bootstrap/tools/go/testmaingen.go | 33 +++++++++++++++++++---------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/go_bootstrap/tools/go/testmaingen.go b/go_bootstrap/tools/go/testmaingen.go
index ec1103b20..79906d0cd 100644
--- a/go_bootstrap/tools/go/testmaingen.go
+++ b/go_bootstrap/tools/go/testmaingen.go
@@ -30,9 +30,11 @@ import (
 	"go/scanner"
 	"go/token"
 	"log"
+	"maps"
 	"os"
 	"path/filepath"
 	"reflect"
+	"slices"
 	"sort"
 	"strings"
 	"text/template"
@@ -131,20 +133,31 @@ func main() {
 	os.Args = loadArgs(os.Args)
 	flag.Parse()
 
-	testFuncs, err := loadTestFuncsFromFiles(pkgImportPath, flag.Args())
-	if err != nil {
-		log.Fatalln("Could not read test files:", err)
-	}
+	coverInfos := make([]coverInfo, 0, len(coverPkgs))
+	for importPath, varToFileMap := range coverPkgs {
+		coverVarMap := make(map[string]*CoverVar, len(varToFileMap))
+		for varName, fileName := range varToFileMap {
+			coverVarMap[fileName] = &CoverVar{
+				File: filepath.Join(importPath, fileName),
+				Var:  varName,
+			}
+		}
 
-	for importPath, vars := range coverPkgs {
-		cover := coverInfo{&Package{importPath}, make(map[string]*CoverVar)}
-		for v, f := range vars {
-			cover.Vars[f] = &CoverVar{File: filepath.Join(importPath, f), Var: v}
+		cover := coverInfo{
+			Package: &Package{ImportPath: importPath},
+			Vars:    coverVarMap,
 		}
-		testFuncs.Cover = append(testFuncs.Cover, cover)
-		testCoverPaths = append(testCoverPaths, importPath)
+		coverInfos = append(coverInfos, cover)
 	}
+
 	testCover = testCoverMode != ""
+	testCoverPaths = append(testCoverPaths, slices.Collect(maps.Keys(coverPkgs))...)
+
+	testFuncs, err := loadTestFuncsFromFiles(pkgImportPath, flag.Args())
+	if err != nil {
+		log.Fatalln("Could not read test files:", err)
+	}
+	testFuncs.Cover = append(testFuncs.Cover, coverInfos...)
 
 	out := os.Stdout
 	if outputFile != "" {