From 55f1d5e3156778020b6e74e13d4b8546ffde9e71 Mon Sep 17 00:00:00 2001
From: Jordan Orelli <jordanorelli@gmail.com>
Date: Sun, 26 Jul 2020 23:14:37 +0000
Subject: [PATCH] define the After interface

---
 examples/tea_test.go | 10 +++++++++-
 test.go              |  8 +++++++-
 tree.go              | 10 ++++++++++
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/examples/tea_test.go b/examples/tea_test.go
index 54edeb7..0b2613e 100644
--- a/examples/tea_test.go
+++ b/examples/tea_test.go
@@ -6,7 +6,7 @@ import (
 	"fmt"
 	"testing"
 
-	"./tea"
+	"github.com/jordanorelli/tea"
 )
 
 type empty struct{}
@@ -25,6 +25,10 @@ func (test *testThingSetup) Run(t *testing.T) {
 	test.Thing = new(Thing)
 }
 
+func (test *testThingSetup) After(t *testing.T) {
+	t.Logf("[%s] testThingSetup after", t.Name())
+}
+
 func (test testThingSetup) String() string { return "thingSetup" }
 
 type setKey struct {
@@ -51,6 +55,10 @@ func (test *setKey) Run(t *testing.T) {
 	}
 }
 
+func (test *setKey) After(t *testing.T) {
+	t.Logf("[%s] setKey after key: %q value: %q", t.Name(), test.key, test.value)
+}
+
 func TestThing(t *testing.T) {
 	root := tea.New(new(testThingSetup))
 
diff --git a/test.go b/test.go
index 7657431..3554b33 100644
--- a/test.go
+++ b/test.go
@@ -9,6 +9,10 @@ type Test interface {
 	Run(*testing.T)
 }
 
+type After interface {
+	After(*testing.T)
+}
+
 func fail(t string, args ...interface{}) Test {
 	return failure{cause: fmt.Errorf(t, args...)}
 }
@@ -17,7 +21,9 @@ type failure struct {
 	cause error
 }
 
-func (f failure) Run(t *testing.T) { t.Error(f.cause.Error()) }
+func (f failure) Run(t *testing.T) {
+	t.Error(f.cause.Error())
+}
 
 // empty is an empty test. It does nothing when run, it's just used as a
 // sentinel value to create notes in the test graph and for ... testing the tea
diff --git a/tree.go b/tree.go
index bf9e126..f8a261f 100644
--- a/tree.go
+++ b/tree.go
@@ -13,6 +13,7 @@ import (
 func Run(t *testing.T, tree *Tree) {
 	t.Run(tree.name, func(t *testing.T) {
 		exec(t, tree)
+		after(t, tree)
 
 		if t.Failed() || t.Skipped() {
 			for _, child := range tree.children {
@@ -47,6 +48,15 @@ func exec(t *testing.T, tree *Tree) *env {
 	return e.save(test)
 }
 
+func after(t *testing.T, tree *Tree) {
+	if a, ok := tree.test.(After); ok {
+		a.After(t)
+	}
+	if tree.parent != nil {
+		after(t, tree.parent)
+	}
+}
+
 // skip skips the provided tree node as well as all of its children.
 func skip(t *testing.T, tree *Tree) {
 	t.Run(tree.name, func(t *testing.T) {