From 7e091dcdb93d999cdeedbbb17132b9f829c336d5 Mon Sep 17 00:00:00 2001
From: Vadym Tishchenko <v.tishchenko@evopay.com.ua>
Date: Tue, 9 Aug 2022 01:08:08 +0300
Subject: [PATCH] Add opportunity to trigger setup/teardown for subtest

---
 suite/interfaces.go | 13 +++++++++++++
 suite/suite.go      | 24 +++++++++++++++++++++++-
 suite/suite_test.go | 29 +++++++++++++++++++++--------
 3 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/suite/interfaces.go b/suite/interfaces.go
index 8b98a8af2..fed037d7f 100644
--- a/suite/interfaces.go
+++ b/suite/interfaces.go
@@ -7,6 +7,7 @@ import "testing"
 type TestingSuite interface {
 	T() *testing.T
 	SetT(*testing.T)
+	SetS(suite TestingSuite)
 }
 
 // SetupAllSuite has a SetupSuite method, which will run before the
@@ -51,3 +52,15 @@ type AfterTest interface {
 type WithStats interface {
 	HandleStats(suiteName string, stats *SuiteInformation)
 }
+
+// SetupSubTest has a SetupSubTest method, which will run before each
+// subtest in the suite.
+type SetupSubTest interface {
+	SetupSubTest()
+}
+
+// TearDownSubTest has a TearDownSubTest method, which will run after
+// each subtest in the suite have been run.
+type TearDownSubTest interface {
+	TearDownSubTest()
+}
diff --git a/suite/suite.go b/suite/suite.go
index 895591878..8b4202d89 100644
--- a/suite/suite.go
+++ b/suite/suite.go
@@ -22,9 +22,13 @@ var matchMethod = flag.String("testify.m", "", "regular expression to select tes
 // retrieving the current *testing.T context.
 type Suite struct {
 	*assert.Assertions
+
 	mu      sync.RWMutex
 	require *require.Assertions
 	t       *testing.T
+
+	// Parent suite to have access to the implemented methods of parent struct
+	s TestingSuite
 }
 
 // T retrieves the current *testing.T context.
@@ -43,6 +47,12 @@ func (suite *Suite) SetT(t *testing.T) {
 	suite.require = require.New(t)
 }
 
+// SetS needs to set the current test suite as parent
+// to get access to the parent methods
+func (suite *Suite) SetS(s TestingSuite) {
+	suite.s = s
+}
+
 // Require returns a require context for suite.
 func (suite *Suite) Require() *require.Assertions {
 	suite.mu.Lock()
@@ -85,7 +95,18 @@ func failOnPanic(t *testing.T, r interface{}) {
 // Provides compatibility with go test pkg -run TestSuite/TestName/SubTestName.
 func (suite *Suite) Run(name string, subtest func()) bool {
 	oldT := suite.T()
-	defer suite.SetT(oldT)
+
+	if setupSubTest, ok := suite.s.(SetupSubTest); ok {
+		setupSubTest.SetupSubTest()
+	}
+
+	defer func() {
+		suite.SetT(oldT)
+		if tearDownSubTest, ok := suite.s.(TearDownSubTest); ok {
+			tearDownSubTest.TearDownSubTest()
+		}
+	}()
+
 	return oldT.Run(name, func(t *testing.T) {
 		suite.SetT(t)
 		subtest()
@@ -98,6 +119,7 @@ func Run(t *testing.T, suite TestingSuite) {
 	defer recoverAndFailOnPanic(t)
 
 	suite.SetT(t)
+	suite.SetS(suite)
 
 	var suiteSetupDone bool
 
diff --git a/suite/suite_test.go b/suite/suite_test.go
index 446029a43..d684f52b9 100644
--- a/suite/suite_test.go
+++ b/suite/suite_test.go
@@ -151,14 +151,16 @@ type SuiteTester struct {
 	Suite
 
 	// Keep counts of how many times each method is run.
-	SetupSuiteRunCount    int
-	TearDownSuiteRunCount int
-	SetupTestRunCount     int
-	TearDownTestRunCount  int
-	TestOneRunCount       int
-	TestTwoRunCount       int
-	TestSubtestRunCount   int
-	NonTestMethodRunCount int
+	SetupSuiteRunCount      int
+	TearDownSuiteRunCount   int
+	SetupTestRunCount       int
+	TearDownTestRunCount    int
+	TestOneRunCount         int
+	TestTwoRunCount         int
+	TestSubtestRunCount     int
+	NonTestMethodRunCount   int
+	SetupSubTestRunCount    int
+	TearDownSubTestRunCount int
 
 	SuiteNameBefore []string
 	TestNameBefore  []string
@@ -255,6 +257,14 @@ func (suite *SuiteTester) TestSubtest() {
 	}
 }
 
+func (suite *SuiteTester) TearDownSubTest() {
+	suite.TearDownSubTestRunCount++
+}
+
+func (suite *SuiteTester) SetupSubTest() {
+	suite.SetupSubTestRunCount++
+}
+
 type SuiteSkipTester struct {
 	// Include our basic suite logic.
 	Suite
@@ -336,6 +346,9 @@ func TestRunSuite(t *testing.T) {
 	assert.Equal(t, suiteTester.TestTwoRunCount, 1)
 	assert.Equal(t, suiteTester.TestSubtestRunCount, 1)
 
+	assert.Equal(t, suiteTester.TearDownSubTestRunCount, 2)
+	assert.Equal(t, suiteTester.SetupSubTestRunCount, 2)
+
 	// Methods that don't match the test method identifier shouldn't
 	// have been run at all.
 	assert.Equal(t, suiteTester.NonTestMethodRunCount, 0)