Ensure that subtest name is set in SetupSubTest/TearDownSubTest

stretchr · Sep 20, 2023 · da0d3c8 · da0d3c8
1 parent 882382d
commit da0d3c8
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 13 deletions.
diff --git a/suite/suite.go b/suite/suite.go
@@ -96,19 +96,20 @@ func failOnPanic(t *testing.T, r interface{}) {
 func (suite *Suite) Run(name string, subtest func()) bool {
 	oldT := suite.T()
 
-	if setupSubTest, ok := suite.s.(SetupSubTest); ok {
-		setupSubTest.SetupSubTest()
-	}
+	return oldT.Run(name, func(t *testing.T) {
+		suite.SetT(t)
 
-	defer func() {
-		suite.SetT(oldT)
-		if tearDownSubTest, ok := suite.s.(TearDownSubTest); ok {
-			tearDownSubTest.TearDownSubTest()
+		if setupSubTest, ok := suite.s.(SetupSubTest); ok {
+			setupSubTest.SetupSubTest()
 		}
-	}()
 
-	return oldT.Run(name, func(t *testing.T) {
-		suite.SetT(t)
+		defer func() {
+			if tearDownSubTest, ok := suite.s.(TearDownSubTest); ok {
+				tearDownSubTest.TearDownSubTest()
+			}
+			suite.SetT(oldT)
+		}()
+
 		subtest()
 	})
 }

diff --git a/suite/suite_test.go b/suite/suite_test.go
@@ -162,6 +162,9 @@ type SuiteTester struct {
 	SetupSubTestRunCount    int
 	TearDownSubTestRunCount int
 
+	SetupSubTestNames    []string
+	TearDownSubTestNames []string
+
 	SuiteNameBefore []string
 	TestNameBefore  []string
 
@@ -170,6 +173,8 @@ type SuiteTester struct {
 
 	TimeBefore []time.Time
 	TimeAfter  []time.Time
+
+	SuiteT *testing.T
 }
 
 // The SetupSuite method will be run by testify once, at the very
@@ -245,23 +250,25 @@ func (suite *SuiteTester) TestSubtest() {
 		{"first"},
 		{"second"},
 	} {
-		suiteT := suite.T()
+		suite.SuiteT = suite.T()
 		suite.Run(t.testName, func() {
 			// We should get a different *testing.T for subtests, so that
 			// go test recognizes them as proper subtests for output formatting
 			// and running individual subtests
 			subTestT := suite.T()
-			suite.NotEqual(subTestT, suiteT)
+			suite.NotEqual(subTestT, suite.SuiteT)
 		})
-		suite.Equal(suiteT, suite.T())
+		suite.Equal(suite.SuiteT, suite.T())
 	}
 }
 
 func (suite *SuiteTester) TearDownSubTest() {
+	suite.TearDownSubTestNames = append(suite.TearDownSubTestNames, suite.T().Name())
 	suite.TearDownSubTestRunCount++
 }
 
 func (suite *SuiteTester) SetupSubTest() {
+	suite.SetupSubTestNames = append(suite.SetupSubTestNames, suite.T().Name())
 	suite.SetupSubTestRunCount++
 }
 
@@ -319,6 +326,12 @@ func TestRunSuite(t *testing.T) {
 	assert.Contains(t, suiteTester.TestNameBefore, "TestSkip")
 	assert.Contains(t, suiteTester.TestNameBefore, "TestSubtest")
 
+	assert.Contains(t, suiteTester.SetupSubTestNames, "TestRunSuite/TestSubtest/first")
+	assert.Contains(t, suiteTester.SetupSubTestNames, "TestRunSuite/TestSubtest/second")
+
+	assert.Contains(t, suiteTester.TearDownSubTestNames, "TestRunSuite/TestSubtest/first")
+	assert.Contains(t, suiteTester.TearDownSubTestNames, "TestRunSuite/TestSubtest/second")
+
 	for _, suiteName := range suiteTester.SuiteNameAfter {
 		assert.Equal(t, "SuiteTester", suiteName)
 	}