fixes issue 290 (#294)

fixtures/deep-exit.go

@@ -1,9 +1,10 @@
 package fixtures
 
 import (
-	"syscall"
 	"log"
 	"os"
+	"syscall"
+	"testing"
 )
 
 func foo0() {
@@ -30,3 +31,8 @@ func bar2() {
 	bar()
 	syscall.Exit(1) // MATCH /calls to syscall.Exit only in main() or init() functions/
 }
+
+func TestMain(m *testing.M) {
+	// must match because this is not a test file
+	os.Exit(m.Run()) // MATCH /calls to os.Exit only in main() or init() functions/
+}

fixtures/deep-exit_test.go

@@ -0,0 +1,11 @@
+package fixtures
+
+import (
+	"os"
+	"testing"
+)
+
+func TestMain(m *testing.M) {
+	// call flag.Parse() here if TestMain uses flags
+	os.Exit(m.Run())
+}

rule/deep-exit.go

@@ -30,7 +30,7 @@ func (r *DeepExitRule) Apply(file *lint.File, _ lint.Arguments) []lint.Failure {
 		},
 	}
 
-	w := lintDeepExit{onFailure, exitFunctions, false}
+	w := lintDeepExit{onFailure, exitFunctions, file.IsTest()}
 	ast.Walk(w, file.AST)
 	return failures
 }
@@ -43,16 +43,15 @@ func (r *DeepExitRule) Name() string {
 type lintDeepExit struct {
 	onFailure     func(lint.Failure)
 	exitFunctions map[string]map[string]bool
-	ignore        bool
+	isTestFile    bool
 }
 
 func (w lintDeepExit) Visit(node ast.Node) ast.Visitor {
-	if stmt, ok := node.(*ast.FuncDecl); ok {
+	if fd, ok := node.(*ast.FuncDecl); ok {
-		w.updateIgnore(stmt)
+		if w.mustIgnore(fd) {
-		return w
+			return nil // skip analysis of this function
 		}
 
-	if w.ignore {
 		return w
 	}
 
@@ -88,7 +87,8 @@ func (w lintDeepExit) Visit(node ast.Node) ast.Visitor {
 	return w
 }
 
-func (w *lintDeepExit) updateIgnore(fd *ast.FuncDecl) {
+func (w *lintDeepExit) mustIgnore(fd *ast.FuncDecl) bool {
 	fn := fd.Name.Name
-	w.ignore = (fn == "init" || fn == "main")
+
+	return fn == "init" || fn == "main" || (w.isTestFile && fn == "TestMain")
 }

test/deep-exit_test.go

@@ -8,4 +8,5 @@ import (
 
 func TestDeepExit(t *testing.T) {
 	testRule(t, "deep-exit", &rule.DeepExitRule{})
+	testRule(t, "deep-exit_test", &rule.DeepExitRule{})
 }