From 1d9d7be28c32d25bfcc557071b5aac285ec06e8a Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Tue, 7 Mar 2017 09:19:02 -0500
Subject: [PATCH] Add schema.Provider.TestReset method

Provider.TestReset resets the internal state of the Provider at the
start of a test. This also adds a MetaReset function field to
schema.Provider, which is called by TestReset and can be used to reset
any other tsated stored in the provider metadata.

This is currently used to reset the internal Context returned by
StopContext between tests, and should  be implemented by a provider if
it stores a Context from a previous test.
---
 helper/schema/provider.go      | 30 ++++++++++++++++++++++++++++++
 helper/schema/provider_test.go | 26 ++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/helper/schema/provider.go b/helper/schema/provider.go
index 5b50d54a1..d52d2f5f0 100644
--- a/helper/schema/provider.go
+++ b/helper/schema/provider.go
@@ -50,8 +50,15 @@ type Provider struct {
 	// See the ConfigureFunc documentation for more information.
 	ConfigureFunc ConfigureFunc
 
+	// MetaReset is called by TestReset to reset any state stored in the meta
+	// interface.  This is especially important if the StopContext is stored by
+	// the provider.
+	MetaReset func() error
+
 	meta interface{}
 
+	// a mutex is required because TestReset can directly repalce the stopCtx
+	stopMu        sync.Mutex
 	stopCtx       context.Context
 	stopCtxCancel context.CancelFunc
 	stopOnce      sync.Once
@@ -124,20 +131,43 @@ func (p *Provider) Stopped() bool {
 // StopCh returns a channel that is closed once the provider is stopped.
 func (p *Provider) StopContext() context.Context {
 	p.stopOnce.Do(p.stopInit)
+
+	p.stopMu.Lock()
+	defer p.stopMu.Unlock()
+
 	return p.stopCtx
 }
 
 func (p *Provider) stopInit() {
+	p.stopMu.Lock()
+	defer p.stopMu.Unlock()
+
 	p.stopCtx, p.stopCtxCancel = context.WithCancel(context.Background())
 }
 
 // Stop implementation of terraform.ResourceProvider interface.
 func (p *Provider) Stop() error {
 	p.stopOnce.Do(p.stopInit)
+
+	p.stopMu.Lock()
+	defer p.stopMu.Unlock()
+
 	p.stopCtxCancel()
 	return nil
 }
 
+// TestReset resets any state stored in the Provider, and will call TestReset
+// on Meta if it implements the TestProvider interface.
+// This may be used to reset the schema.Provider at the start of a test, and is
+// automatically called by resource.Test.
+func (p *Provider) TestReset() error {
+	p.stopInit()
+	if p.MetaReset != nil {
+		return p.MetaReset()
+	}
+	return nil
+}
+
 // Input implementation of terraform.ResourceProvider interface.
 func (p *Provider) Input(
 	input terraform.UIInput,
diff --git a/helper/schema/provider_test.go b/helper/schema/provider_test.go
index ed5918844..5b06c5e57 100644
--- a/helper/schema/provider_test.go
+++ b/helper/schema/provider_test.go
@@ -381,3 +381,29 @@ func TestProviderStop_stopFirst(t *testing.T) {
 		t.Fatal("should be stopped")
 	}
 }
+
+func TestProviderReset(t *testing.T) {
+	var p Provider
+	stopCtx := p.StopContext()
+	p.MetaReset = func() error {
+		stopCtx = p.StopContext()
+		return nil
+	}
+
+	// cancel the current context
+	p.Stop()
+
+	if err := p.TestReset(); err != nil {
+		t.Fatal(err)
+	}
+
+	// the first context should have been replaced
+	if err := stopCtx.Err(); err != nil {
+		t.Fatal(err)
+	}
+
+	// we should not get a canceled context here either
+	if err := p.StopContext().Err(); err != nil {
+		t.Fatal(err)
+	}
+}