fixes issue 290 (#294)

Author: chavacava

fixtures/deep-exit.go

@@ -1,32 +1,38 @@
 package fixtures
 
 import (
-	"syscall"
 	"log"
 	"os"
+	"syscall"
+	"testing"
 )
 
 func foo0() {
-    os.Exit(1) // MATCH /calls to os.Exit only in main() or init() functions/
+	os.Exit(1) // MATCH /calls to os.Exit only in main() or init() functions/
 }
 
 func init() {
-    log.Fatal("v ...interface{}")
+	log.Fatal("v ...interface{}")
 }
 
 func foo() {
-    log.Fatalf(1) // MATCH /calls to log.Fatalf only in main() or init() functions/
+	log.Fatalf(1) // MATCH /calls to log.Fatalf only in main() or init() functions/
 }
 
 func main() {
-    log.Fatalln("v ...interface{}")
+	log.Fatalln("v ...interface{}")
 }
 
 func bar() {
-    log.Fatal(1) // MATCH /calls to log.Fatal only in main() or init() functions/
+	log.Fatal(1) // MATCH /calls to log.Fatal only in main() or init() functions/
 }
 
 func bar2() {
-    bar()
-    syscall.Exit(1) // MATCH /calls to syscall.Exit only in main() or init() functions/
+	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 {
-		w.updateIgnore(stmt)
-		return w
-	}
+	if fd, ok := node.(*ast.FuncDecl); ok {
+		if w.mustIgnore(fd) {
+			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{})
 }