1 files changed, 65 insertions, 11 deletions

diff --git a/tools/golangci-lint/gitaly/testhelper_run.go b/tools/golangci-lint/gitaly/testhelper_run.go
index aa9226c02..116ab4162 100644
--- a/tools/golangci-lint/gitaly/testhelper_run.go
+++ b/tools/golangci-lint/gitaly/testhelper_run.go
@@ -17,6 +17,17 @@ type testhelperRunAnalyzerSettings struct {
 	IncludedFunctions []string `mapstructure:"included-functions"`
 }
 
+// testmainFact is used to report if a package has defined a `TestMain` function.
+type testmainFact struct {
+	HasTestMain bool
+}
+
+// AFact is used to satisfy the `Fact` interface.
+func (*testmainFact) AFact() {}
+
+// String returns the message expected by tests when a testmainFact is exported.
+func (*testmainFact) String() string { return "package has TestMain" }
+
 var toolPrefixPattern = regexp.MustCompile(`^gitlab.com/gitlab-org/gitaly(/v\d{2})?/tools`)
 
 // newTesthelperRunAnalyzer returns an analyzer to detect if a package that has tests does
@@ -25,23 +36,17 @@ var toolPrefixPattern = regexp.MustCompile(`^gitlab.com/gitlab-org/gitaly(/v\d{2
 // https://gitlab.com/gitlab-org/gitaly/-/blob/master/STYLE.md?ref_type=heads#common-setup
 func newTesthelperRunAnalyzer(settings *testhelperRunAnalyzerSettings) *analysis.Analyzer {
 	return &analysis.Analyzer{
-		Name: testhelperRunAnalyzerName,
-		Doc:  `TestMain must be present and call testhelper.Run()`,
-		Run:  runTesthelperRunAnalyzer(settings.IncludedFunctions),
+		Name:      testhelperRunAnalyzerName,
+		Doc:       `TestMain must be present and call testhelper.Run()`,
+		Run:       runTesthelperRunAnalyzer(settings.IncludedFunctions),
+		FactTypes: []analysis.Fact{&testmainFact{}},
 	}
 }
 
 func runTesthelperRunAnalyzer(rules []string) func(*analysis.Pass) (interface{}, error) {
 	return func(pass *analysis.Pass) (interface{}, error) {
 		var hasTestMain, hasTests bool
-
-		// Blackbox test packages ending with `_test` are considered to be
-		// part of the primary package for compilation, but are scanned in a
-		// separate pass by the analyzer. The primary and test packages cannot
-		// both define `TestMain`. Only require `TestMain` in the primary package.
-		if strings.HasSuffix(pass.Pkg.Name(), "_test") {
-			return nil, nil
-		}
+		var fact testmainFact
 
 		// Don't lint tools, they can't import `testhelper`.
 		if toolPrefixPattern.MatchString(pass.Pkg.Path()) {
@@ -53,12 +58,23 @@ func runTesthelperRunAnalyzer(rules []string) func(*analysis.Pass) (interface{},
 				break
 			}
 
+			// Blackbox test packages ending with `_test` are considered to be
+			// part of the primary package for compilation, but are scanned in a
+			// separate pass by the analyzer. The primary and test packages cannot
+			// both define `TestMain`.
+			if isTestPkg(pass) && primaryPkgHasTestMain(pass, file) {
+				// Primary package has already defined `TestMain`, no need to check
+				// `_test` package.
+				break
+			}
+
 			ast.Inspect(file, func(node ast.Node) bool {
 				if decl, ok := node.(*ast.FuncDecl); ok {
 					declName := decl.Name.Name
 
 					if declName == "TestMain" {
 						hasTestMain = true
+						fact.HasTestMain = true
 
 						analyzeTestMain(pass, decl, rules)
 						analyzeFilename(pass, file, decl)
@@ -94,6 +110,9 @@ func runTesthelperRunAnalyzer(rules []string) func(*analysis.Pass) (interface{},
 			})
 		}
 
+		if hasTestMain {
+			pass.ExportPackageFact(&fact)
+		}
 		return nil, nil
 	}
 }
@@ -134,3 +153,38 @@ func analyzeTestMain(pass *analysis.Pass, decl *ast.FuncDecl, rules []string) {
 		})
 	}
 }
+
+func primaryPkgHasTestMain(pass *analysis.Pass, file *ast.File) bool {
+	var primaryPkg *types.Package
+
+	ast.Inspect(file, func(node ast.Node) bool {
+		if spec, ok := node.(*ast.ImportSpec); ok {
+			obj, ok := pass.TypesInfo.Implicits[spec]
+			if !ok {
+				obj = pass.TypesInfo.Defs[spec.Name]
+			}
+			importedPkg := obj.(*types.PkgName).Imported()
+
+			if importedPkg.Path() == strings.TrimSuffix(pass.Pkg.Path(), "_test") {
+				primaryPkg = importedPkg
+			}
+		}
+
+		return true
+	})
+
+	if primaryPkg == nil {
+		return false
+	}
+
+	var primaryFact testmainFact
+	if !pass.ImportPackageFact(primaryPkg, &primaryFact) {
+		return false
+	}
+
+	return primaryFact.HasTestMain
+}
+
+func isTestPkg(pass *analysis.Pass) bool {
+	return strings.HasSuffix(pass.Pkg.Name(), "_test")
+}