From ce0137fa6a71702fcac984cb87f6460f7b61d7c4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Sun, 16 Feb 2020 17:16:49 +0000
Subject: [PATCH] don't break TestMain funcs

Important for 'garble test', if a package uses one.
---
 main.go                   |  2 +-
 testdata/scripts/test.txt | 11 +++++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/main.go b/main.go
index 1316955..5b0d3af 100644
--- a/main.go
+++ b/main.go
@@ -363,7 +363,7 @@ func transformGo(node ast.Node, info *types.Info) ast.Node {
 					return true // might implement an interface
 				}
 				switch node.Name {
-				case "main", "init":
+				case "main", "init", "TestMain":
 					return true // don't break them
 				}
 				if strings.HasPrefix(node.Name, "Test") && isTestSignature(sign) {
diff --git a/testdata/scripts/test.txt b/testdata/scripts/test.txt
index 9b0482b..bc148c9 100644
--- a/testdata/scripts/test.txt
+++ b/testdata/scripts/test.txt
@@ -45,3 +45,14 @@ func TestSeparateFoo(t *testing.T) {
 		t.FailNow()
 	}
 }
+-- main_test.go --
+package bar
+
+import (
+	"os"
+	"testing"
+)
+
+func TestMain(m *testing.M) {
+	os.Exit(m.Run())
+}