From e8910bb33558a8e2742c40f644429545222bea7b Mon Sep 17 00:00:00 2001
From: Esdras Beleza <github@contato.esdrasbeleza.com>
Date: Tue, 3 Mar 2020 19:48:03 +0000
Subject: [PATCH] Add stats to suites and tests

---
 suite/interfaces.go |  9 ++++++++-
 suite/stats.go      | 38 +++++++++++++++++++++++++++++++++++++
 suite/suite.go      | 46 +++++++++++++++++++++++++++++++++++++++++++--
 suite/suite_test.go | 30 +++++++++++++++++++++++++++++
 4 files changed, 120 insertions(+), 3 deletions(-)
 create mode 100644 suite/stats.go

diff --git a/suite/interfaces.go b/suite/interfaces.go
index b37cb04..303459d 100644
--- a/suite/interfaces.go
+++ b/suite/interfaces.go
@@ -1,6 +1,8 @@
 package suite
 
-import "testing"
+import (
+	"testing"
+)
 
 // TestingSuite can store and return the current *testing.T context
 // generated by 'go test'.
@@ -44,3 +46,8 @@ type BeforeTest interface {
 type AfterTest interface {
 	AfterTest(suiteName, testName string)
 }
+
+// WithStats is a suite which measures some stats for its tests.
+type WithStats interface {
+	HandleStats(suiteName string, stats *SuiteInformation)
+}
diff --git a/suite/stats.go b/suite/stats.go
new file mode 100644
index 0000000..df52b10
--- /dev/null
+++ b/suite/stats.go
@@ -0,0 +1,38 @@
+package suite
+
+import "time"
+
+// SuiteInformation stats stores stats for the whole suite execution.
+type SuiteInformation struct {
+	Start, End time.Time
+	Passed     bool
+	TestStats  map[string]*TestInformation
+}
+
+// TestInformation stores information about the execution of each test.
+type TestInformation struct {
+	TestName   string
+	Start, End time.Time
+	Passed     bool
+}
+
+func newSuiteInformation() *SuiteInformation {
+	testStats := make(map[string]*TestInformation)
+
+	return &SuiteInformation{
+		TestStats: testStats,
+		Passed:    true,
+	}
+}
+
+func (s SuiteInformation) start(testName string) {
+	s.TestStats[testName] = &TestInformation{
+		TestName: testName,
+		Start:    time.Now(),
+	}
+}
+
+func (s SuiteInformation) end(testName string, passed bool) {
+	s.TestStats[testName].End = time.Now()
+	s.TestStats[testName].Passed = passed
+}
diff --git a/suite/suite.go b/suite/suite.go
index 6195301..26452dd 100644
--- a/suite/suite.go
+++ b/suite/suite.go
@@ -9,6 +9,7 @@ import (
 	"runtime/debug"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -87,6 +88,14 @@ func Run(t *testing.T, suite TestingSuite) {
 
 	suiteSetupDone := false
 
+	var (
+		stats *SuiteInformation
+	)
+
+	if _, measureStats := suite.(WithStats); measureStats {
+		stats = newSuiteInformation()
+	}
+
 	methodFinder := reflect.TypeOf(suite)
 	tests := []testing.InternalTest{}
 	for index := 0; index < methodFinder.NumMethod(); index++ {
@@ -96,21 +105,36 @@ func Run(t *testing.T, suite TestingSuite) {
 			fmt.Fprintf(os.Stderr, "testify: invalid regexp for -m: %s\n", err)
 			os.Exit(1)
 		}
+
 		if !ok {
 			continue
 		}
+
+		suiteName := methodFinder.Elem().Name()
+
 		if !suiteSetupDone {
+			if stats != nil {
+				stats.Start = time.Now()
+			}
+
 			if setupAllSuite, ok := suite.(SetupAllSuite); ok {
 				setupAllSuite.SetupSuite()
 			}
+
 			defer func() {
 				if tearDownAllSuite, ok := suite.(TearDownAllSuite); ok {
 					testsSync.Wait()
 					tearDownAllSuite.TearDownSuite()
 				}
+
+				if suiteWithStats, measureStats := suite.(WithStats); measureStats {
+					stats.End = time.Now()
+					suiteWithStats.HandleStats(suiteName, stats)
+				}
 			}()
 			suiteSetupDone = true
 		}
+
 		test := testing.InternalTest{
 			Name: method.Name,
 			F: func(t *testing.T) {
@@ -122,16 +146,34 @@ func Run(t *testing.T, suite TestingSuite) {
 				if setupTestSuite, ok := suite.(SetupTestSuite); ok {
 					setupTestSuite.SetupTest()
 				}
+
 				if beforeTestSuite, ok := suite.(BeforeTest); ok {
 					beforeTestSuite.BeforeTest(methodFinder.Elem().Name(), method.Name)
 				}
+
+				if stats != nil {
+					stats.start(method.Name)
+				}
+
 				defer func() {
-					if afterTestSuite, ok := suite.(AfterTest); ok {
-						afterTestSuite.AfterTest(methodFinder.Elem().Name(), method.Name)
+					if stats != nil {
+						passed := !t.Failed()
+
+						stats.end(method.Name, passed)
+
+						if !passed {
+							stats.Passed = false
+						}
 					}
+
+					if afterTestSuite, ok := suite.(AfterTest); ok {
+						afterTestSuite.AfterTest(suiteName, method.Name)
+					}
+
 					if tearDownTestSuite, ok := suite.(TearDownTestSuite); ok {
 						tearDownTestSuite.TearDownTest()
 					}
+
 					suite.SetT(parentT)
 				}()
 				method.Func.Call([]reflect.Value{reflect.ValueOf(suite)})
diff --git a/suite/suite_test.go b/suite/suite_test.go
index 5d757dd..9da9223 100644
--- a/suite/suite_test.go
+++ b/suite/suite_test.go
@@ -482,3 +482,33 @@ func (s *CallOrderSuite) Test_A() {
 func (s *CallOrderSuite) Test_B() {
 	s.call("Test B")
 }
+
+type suiteWithStats struct {
+	Suite
+	wasCalled bool
+	stats     *SuiteInformation
+}
+
+func (s *suiteWithStats) HandleStats(suiteName string, stats *SuiteInformation) {
+	s.wasCalled = true
+	s.stats = stats
+}
+
+func (s *suiteWithStats) TestSomething() {
+	s.Equal(1, 1)
+}
+
+func TestSuiteWithStats(t *testing.T) {
+	suiteWithStats := new(suiteWithStats)
+	Run(t, suiteWithStats)
+
+	assert.True(t, suiteWithStats.wasCalled)
+	assert.NotZero(t, suiteWithStats.stats.Start)
+	assert.NotZero(t, suiteWithStats.stats.End)
+	assert.True(t, suiteWithStats.stats.Passed)
+
+	testStats := suiteWithStats.stats.TestStats["TestSomething"]
+	assert.NotZero(t, testStats.Start)
+	assert.NotZero(t, testStats.End)
+	assert.True(t, testStats.Passed)
+}