From b915aa36bc999e07269268d07a52d03ffb17bf1e Mon Sep 17 00:00:00 2001
From: Adam Luzsi <adamluzsi@gmail.com>
Date: Tue, 17 Sep 2024 20:44:03 +0200
Subject: [PATCH] add support for Spec#AfterAll

`Spec#AfterAll` is perfect for cleaning up shared states across tests (like a DB connection).
It's also useful when you need to make assertions that should only happen after all the tests have finished running.
---
 Spec.go       | 27 +++++++++++++++
 hooks.go      | 27 +++++++++++++--
 hooks_test.go | 94 +++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 146 insertions(+), 2 deletions(-)

diff --git a/Spec.go b/Spec.go
index c9658cf..cd1cc78 100644
--- a/Spec.go
+++ b/Spec.go
@@ -90,6 +90,7 @@ type Spec struct {
 	hooks struct {
 		Around    []hook
 		BeforeAll []hookOnce
+		AfterAll  []hookOnce
 	}
 
 	defs []func(*Spec)
@@ -602,6 +603,7 @@ func (spec *Spec) Finish() {
 		spec.orderer.Order(tests)
 		td := &teardown.Teardown{}
 		defer spec.documentResults()
+		defer spec.runAfterAll()
 		defer td.Finish()
 		for _, tc := range tests {
 			tc()
@@ -609,6 +611,31 @@ func (spec *Spec) Finish() {
 	})
 }
 
+func (spec *Spec) runAfterAll() {
+	helper(spec.testingTB).Helper()
+	if spec.testingTB == nil {
+		return
+	}
+	if spec.parent != nil {
+		return
+	}
+	if spec.isSuite() {
+		return
+	}
+	spec.visitAll(func(s *Spec) {
+		for _, h := range s.hooks.AfterAll {
+			h.DoOnce(spec.testingTB)
+		}
+	})
+}
+
+func (spec *Spec) visitAll(fn func(*Spec)) {
+	fn(spec)
+	for _, child := range spec.children {
+		child.visitAll(fn)
+	}
+}
+
 func (spec *Spec) documentResults() {
 	if spec.testingTB == nil {
 		return
diff --git a/hooks.go b/hooks.go
index 6537d14..09aa980 100644
--- a/hooks.go
+++ b/hooks.go
@@ -42,8 +42,6 @@ func (spec *Spec) Before(beforeBlock tBlock) {
 // The received *testing.T object is the same as the Then block *testing.T object
 // This hook applied to this scope and anything that is nested from here.
 // All setup block is stackable.
-//
-// DEPRECATED: use Spec.Before with T.Cleanup or Spec.Before with T.Defer instead
 func (spec *Spec) After(afterBlock tBlock) {
 	helper(spec.testingTB).Helper()
 	spec.Around(func(t *T) func() {
@@ -99,3 +97,28 @@ func (spec *Spec) BeforeAll(blk func(tb testing.TB)) {
 		spec.hooks.BeforeAll = append(spec.hooks.BeforeAll, h)
 	})
 }
+
+func (spec *Spec) AfterAll(blk func(tb testing.TB)) {
+	helper(spec.testingTB).Helper()
+	frame, _ := caller.GetFrame()
+	spec.modify(func(spec *Spec) {
+		helper(spec.testingTB).Helper()
+
+		if spec.immutable {
+			spec.testingTB.Fatal(hookWarning)
+		}
+
+		var onCall sync.Once
+		var beforeAll = func(tb testing.TB) {
+			onCall.Do(func() { blk(tb) })
+		}
+
+		h := hookOnce{
+			DoOnce: beforeAll,
+			Block:  blk,
+			Frame:  frame,
+		}
+
+		spec.hooks.AfterAll = append(spec.hooks.AfterAll, h)
+	})
+}
diff --git a/hooks_test.go b/hooks_test.go
index 09cb377..3b5a7b4 100644
--- a/hooks_test.go
+++ b/hooks_test.go
@@ -92,3 +92,97 @@ func TestSpec_BeforeAll_failIfDefinedAfterTestCases(t *testing.T) {
 	})
 	assert.Must(t).True(stub.IsFailed)
 }
+
+func ExampleSpec_AfterAll() {
+	s := testcase.NewSpec(nil)
+	s.AfterAll(func(tb testing.TB) {
+		// do something after all the test finished running
+	})
+	s.Test("this test will run before the AfterAll hook", func(t *testcase.T) {})
+}
+
+func TestSpec_AfterAll(t *testing.T) {
+	stub := &doubles.TB{}
+	var order []string
+	sandbox.Run(func() {
+		s := testcase.NewSpec(stub)
+		s.HasSideEffect()
+		s.AfterAll(func(tb testing.TB) {
+			order = append(order, "AfterAll")
+		})
+		s.Test(``, func(t *testcase.T) {
+			order = append(order, "Test")
+		})
+		s.Test(``, func(t *testcase.T) {
+			order = append(order, "Test")
+		})
+		s.Finish()
+	})
+	assert.Must(t).False(stub.IsFailed)
+	assert.Equal(t, []string{"Test", "Test", "AfterAll"}, order,
+		`expected to only run once (single "AfterAll" in the order array)`,
+		`and it should have run in order (After all the "Test")`,
+	)
+}
+
+func TestSpec_AfterAll_nested(t *testing.T) {
+	stub := &doubles.TB{}
+	var order []string
+	sandbox.Run(func() {
+		s := testcase.NewSpec(stub)
+		s.HasSideEffect()
+		s.AfterAll(func(tb testing.TB) {
+			order = append(order, "AfterAll")
+		})
+		s.Context(``, func(s *testcase.Spec) {
+			s.AfterAll(func(tb testing.TB) {
+				order = append(order, "AfterAll")
+			})
+			s.Test(``, func(t *testcase.T) {
+				order = append(order, "Test")
+			})
+		})
+		s.Test(``, func(t *testcase.T) {
+			order = append(order, "Test")
+		})
+		s.Finish()
+	})
+	assert.Must(t).False(stub.IsFailed)
+	assert.Equal(t, []string{"Test", "Test", "AfterAll", "AfterAll"}, order)
+}
+
+func TestSpec_AfterAll_suite(t *testing.T) {
+	stub := &doubles.TB{}
+	var order []string
+	sandbox.Run(func() {
+		suiteSpec1 := testcase.NewSpec(nil)
+		suiteSpec1.HasSideEffect()
+		suiteSpec1.AfterAll(func(tb testing.TB) {
+			order = append(order, "AfterAll")
+		})
+		suiteSpec1.Test(``, func(t *testcase.T) {
+			order = append(order, "Test")
+		})
+		suite1 := suiteSpec1.AsSuite("suite")
+		suiteSpec2 := testcase.NewSpec(nil)
+		suiteSpec2.Context("", suite1.Spec)
+
+		ss := testcase.NewSpec(stub)
+		ss.Context("", suiteSpec2.Spec)
+		ss.Finish()
+	})
+	assert.Must(t).False(stub.IsFailed)
+	assert.Equal(t, []string{"Test", "AfterAll"}, order, "expected to only run once, in the real spec execution")
+}
+
+func TestSpec_AfterAll_failIfDefinedAfterTestCases(t *testing.T) {
+	stub := &doubles.TB{}
+	sandbox.Run(func() {
+		s := testcase.NewSpec(stub)
+		s.Test(``, func(t *testcase.T) {})
+		s.AfterAll(func(tb testing.TB) {})
+		s.Test(``, func(t *testcase.T) {})
+		s.Finish()
+	})
+	assert.Must(t).True(stub.IsFailed)
+}