From 3d3789920def6f2b46fcbd1741cb46a2a7ec3e6f Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Thu, 21 Aug 2014 22:15:47 -0700
Subject: [PATCH] helper/schema: can set conninfo

---
 helper/schema/resource_data.go      | 20 ++++++++++++++++++++
 helper/schema/resource_data_test.go | 16 ++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/helper/schema/resource_data.go b/helper/schema/resource_data.go
index ccb1183f9..eb8eaaba3 100644
--- a/helper/schema/resource_data.go
+++ b/helper/schema/resource_data.go
@@ -91,6 +91,19 @@ func (d *ResourceData) Id() string {
 	return result
 }
 
+// ConnInfo returns the connection info for this resource.
+func (d *ResourceData) ConnInfo() map[string]string {
+	if d.newState != nil {
+		return d.newState.ConnInfo
+	}
+
+	if d.state != nil {
+		return d.state.ConnInfo
+	}
+
+	return nil
+}
+
 // Dependencies returns the dependencies in this state.
 func (d *ResourceData) Dependencies() []terraform.ResourceDependency {
 	if d.newState != nil {
@@ -111,6 +124,12 @@ func (d *ResourceData) SetId(v string) {
 	d.newState.ID = v
 }
 
+// SetConnInfo sets the connection info for a resource.
+func (d *ResourceData) SetConnInfo(v map[string]string) {
+	d.once.Do(d.init)
+	d.newState.ConnInfo = v
+}
+
 // SetDependencies sets the dependencies of a resource.
 func (d *ResourceData) SetDependencies(ds []terraform.ResourceDependency) {
 	d.once.Do(d.init)
@@ -123,6 +142,7 @@ func (d *ResourceData) State() *terraform.ResourceState {
 	var result terraform.ResourceState
 	result.ID = d.Id()
 	result.Attributes = d.stateObject("", d.schema)
+	result.ConnInfo = d.ConnInfo()
 	result.Dependencies = d.Dependencies()
 
 	if v := d.Id(); v != "" {
diff --git a/helper/schema/resource_data_test.go b/helper/schema/resource_data_test.go
index 072e8d7d8..3fee496c3 100644
--- a/helper/schema/resource_data_test.go
+++ b/helper/schema/resource_data_test.go
@@ -1531,6 +1531,22 @@ func TestResourceDataState(t *testing.T) {
 	}
 }
 
+func TestResourceDataSetConnInfo(t *testing.T) {
+	d := &ResourceData{}
+	d.SetConnInfo(map[string]string{
+		"foo": "bar",
+	})
+
+	expected := map[string]string{
+		"foo": "bar",
+	}
+
+	actual := d.State()
+	if !reflect.DeepEqual(actual.ConnInfo, expected) {
+		t.Fatalf("bad: %#v", actual)
+	}
+}
+
 func TestResourceDataSetDependencies(t *testing.T) {
 	d := &ResourceData{}
 	d.SetDependencies([]terraform.ResourceDependency{