Manage time resetting better

mockoidc.go

@@ -171,9 +171,17 @@ func (m *MockOIDC) Now() time.Time {
 	return NowFunc().Add(m.fastForward)
 }
 
-// Synchronize sets the jwt.TimeFunc to our mutated view of time
-func (m *MockOIDC) Synchronize() {
+// TimeReset is a function that resets time
+type TimeReset func()
+
+// Synchronize sets the jwt.TimeFunc to our mutated view of time.
+// It returns a func that can reset it to its original state.
+func (m *MockOIDC) Synchronize() TimeReset {
+	original := jwt.TimeFunc
+
 	jwt.TimeFunc = m.Now
+
+	return func() { jwt.TimeFunc = original }
 }
 
 func (m *MockOIDC) Addr() string {

mockoidc_test.go

@@ -8,7 +8,6 @@ import (
 	"net/http"
 	"net/url"
 	"testing"
-	"time"
 
 	"github.com/dgrijalva/jwt-go"
 	"github.com/oauth2-proxy/mockoidc/v1"
@@ -28,10 +27,8 @@ func TestRun(t *testing.T) {
 	defer m.Shutdown()
 
 	// Override jwt.TimeFunc with our timer
-	m.Synchronize()
-	defer func() {
-		jwt.TimeFunc = time.Now
-	}()
+	reset := m.Synchronize()
+	defer reset()
 
 	// ************************************************************************
 	// Stage 0: Get Discovery documents