From d20ac9444aef0fe1752865d46d69c45f8ec6581f Mon Sep 17 00:00:00 2001
From: NiseVoid <nisevoid@gmail.com>
Date: Tue, 12 May 2020 18:00:07 +0200
Subject: [PATCH] Use Cleanup to remove old entries in ts map

---
 assert.go | 4 ++++
 type.go   | 1 +
 2 files changed, 5 insertions(+)

diff --git a/assert.go b/assert.go
index 1a3166b..bf26a02 100644
--- a/assert.go
+++ b/assert.go
@@ -24,6 +24,10 @@ func New(t T) Assert {
 	f := runtime.FuncForPC(reflect.ValueOf(a).Pointer())
 	ts[f] = t
 
+	t.Cleanup(func() {
+		delete(ts, f)
+	})
+
 	return a
 }
 
diff --git a/type.go b/type.go
index 1c9a1fd..2a4b72a 100644
--- a/type.go
+++ b/type.go
@@ -4,4 +4,5 @@ package assert
 type T interface {
 	Error(...interface{})
 	Helper()
+	Cleanup(func())
 }