From 42a2b64695483f309b6bd04e7f57dbb13266cc05 Mon Sep 17 00:00:00 2001
From: Fatih Arslan <ftharsln@gmail.com>
Date: Wed, 3 Apr 2024 15:10:44 +0300
Subject: [PATCH] Add '--only-tests' support

Signed-off-by: Fatih Arslan <ftharsln@gmail.com>
---
 faillint/faillint.go                           | 18 +++++++++++++++++-
 faillint/faillint_test.go                      | 18 +++++++++++++++++-
 .../testdata/src/sleepintest/sleepintest.go    |  9 +++++++++
 .../src/sleepintest_err/sleepintest_test.go    | 10 ++++++++++
 4 files changed, 53 insertions(+), 2 deletions(-)
 create mode 100644 faillint/testdata/src/sleepintest/sleepintest.go
 create mode 100644 faillint/testdata/src/sleepintest_err/sleepintest_test.go

diff --git a/faillint/faillint.go b/faillint/faillint.go
index 52ff408..2b33d13 100644
--- a/faillint/faillint.go
+++ b/faillint/faillint.go
@@ -3,6 +3,7 @@
 package faillint
 
 import (
+	"errors"
 	"fmt"
 	"go/ast"
 	"go/token"
@@ -62,6 +63,7 @@ type path struct {
 type faillint struct {
 	paths       string // -paths flag
 	ignoretests bool   // -ignore-tests flag
+	onlyTests   bool   // -only-tests flag
 }
 
 // NewAnalyzer create a faillint analyzer.
@@ -69,6 +71,7 @@ func NewAnalyzer() *analysis.Analyzer {
 	f := faillint{
 		paths:       "",
 		ignoretests: false,
+		onlyTests:   false,
 	}
 	a := &analysis.Analyzer{
 		Name:             "faillint",
@@ -90,7 +93,9 @@ Fail on the usage of prometheus.DefaultGatherer and prometheus.MustRegister
 
 Fail on the usage of errors, golang.org/x/net and all sub packages under golang.org/x/net
   -paths errors,golang.org/x/net/...`)
+
 	a.Flags.BoolVar(&f.ignoretests, "ignore-tests", false, "ignore all _test.go files")
+	a.Flags.BoolVar(&f.onlyTests, "only-tests", false, "include only _test.go files")
 	return a
 }
 
@@ -111,6 +116,10 @@ func (f *faillint) run(pass *analysis.Pass) (interface{}, error) {
 		return nil, nil
 	}
 
+	if f.ignoretests && f.onlyTests {
+		return nil, errors.New("--ignore-tests and --only-tests flags cannot be used together")
+	}
+
 	for _, file := range pass.Files {
 		filename := pass.Fset.File(file.Package).Name()
 		isGenerated, err := generated.ParseFile(filename)
@@ -122,9 +131,16 @@ func (f *faillint) run(pass *analysis.Pass) (interface{}, error) {
 			continue
 		}
 
-		if f.ignoretests && strings.Contains(pass.Fset.File(file.Package).Name(), "_test.go") {
+		isTestFile := strings.Contains(pass.Fset.File(file.Package).Name(), "_test.go")
+
+		if f.ignoretests && isTestFile {
+			continue
+		}
+
+		if f.onlyTests && !isTestFile {
 			continue
 		}
+
 		if anyHasDirective(pass, file.Comments, fileIgnoreKey) {
 			continue
 		}
diff --git a/faillint/faillint_test.go b/faillint/faillint_test.go
index 954da8c..441fcec 100644
--- a/faillint/faillint_test.go
+++ b/faillint/faillint_test.go
@@ -188,8 +188,21 @@ func TestRun(t *testing.T) {
 		dir   string
 		paths string
 
-		ignoreTestFiles bool
+		ignoreTestFiles   bool
+		onlyTestFunctions bool
 	}{
+		{
+			name:              "sleep in a function which is not a test",
+			dir:               "sleepintest",
+			paths:             "time.{Sleep}",
+			onlyTestFunctions: true,
+		},
+		{
+			name:              "sleep in a function which is a test",
+			dir:               "sleepintest_err",
+			paths:             "time.{Sleep}",
+			onlyTestFunctions: true,
+		},
 		{
 			name:  "unwanted errors package present",
 			dir:   "a",
@@ -320,6 +333,9 @@ func TestRun(t *testing.T) {
 			if tcase.ignoreTestFiles {
 				f.Flags.Set("ignore-tests", "true")
 			}
+			if tcase.onlyTestFunctions {
+				f.Flags.Set("only-tests", "true")
+			}
 
 			// No assertion on result is required as 'analysistest' is for that.
 			// All expected diagnosis should be specified by comment in affected file starting with `// want`.
diff --git a/faillint/testdata/src/sleepintest/sleepintest.go b/faillint/testdata/src/sleepintest/sleepintest.go
new file mode 100644
index 0000000..0c55e5d
--- /dev/null
+++ b/faillint/testdata/src/sleepintest/sleepintest.go
@@ -0,0 +1,9 @@
+package sleepintest
+
+import (
+	"time"
+)
+
+func fooTest() {
+	time.Sleep(1 * time.Second)
+}
diff --git a/faillint/testdata/src/sleepintest_err/sleepintest_test.go b/faillint/testdata/src/sleepintest_err/sleepintest_test.go
new file mode 100644
index 0000000..b00ee1a
--- /dev/null
+++ b/faillint/testdata/src/sleepintest_err/sleepintest_test.go
@@ -0,0 +1,10 @@
+package sleepintest
+
+import (
+	"testing"
+	"time"
+)
+
+func TestFoo(t *testing.T) {
+	time.Sleep(1 * time.Second) // want `declaration "Sleep" from package "time" shouldn't be used`
+}