From ca18e971d13ea5432faf4adbee7e80d93c45d39b Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Wed, 20 Aug 2014 18:11:14 -0700
Subject: [PATCH] helper/schema: can set sets

---
 helper/schema/resource_data.go      | 26 +++++++-
 helper/schema/resource_data_test.go | 98 +++++++++++++++++++++++++++++
 2 files changed, 123 insertions(+), 1 deletion(-)

diff --git a/helper/schema/resource_data.go b/helper/schema/resource_data.go
index 5a2f0b710..4b9217305 100644
--- a/helper/schema/resource_data.go
+++ b/helper/schema/resource_data.go
@@ -480,8 +480,16 @@ func (d *ResourceData) set(
 		return d.setList(k, parts, schema, value)
 	case TypeMap:
 		return d.setMapValue(k, parts, schema, value)
-	default:
+	case TypeSet:
+		return d.setSet(k, parts, schema, value)
+	case TypeBool:
+		fallthrough
+	case TypeInt:
+		fallthrough
+	case TypeString:
 		return d.setPrimitive(k, schema, value)
+	default:
+		panic(fmt.Sprintf("%s: unknown type %s", k, schema.Type))
 	}
 }
 
@@ -661,6 +669,22 @@ func (d *ResourceData) setPrimitive(
 	return nil
 }
 
+func (d *ResourceData) setSet(
+	k string,
+	parts []string,
+	schema *Schema,
+	value interface{}) error {
+	if len(parts) > 0 {
+		return fmt.Errorf("%s: can only set the full set, not elements", k)
+	}
+
+	if s, ok := value.(*Set); ok {
+		value = s.List()
+	}
+
+	return d.setList(k, nil, schema, value)
+}
+
 func (d *ResourceData) stateList(
 	prefix string,
 	schema *Schema) map[string]string {
diff --git a/helper/schema/resource_data_test.go b/helper/schema/resource_data_test.go
index b311dd78c..072e8d7d8 100644
--- a/helper/schema/resource_data_test.go
+++ b/helper/schema/resource_data_test.go
@@ -1055,6 +1055,101 @@ func TestResourceDataSet(t *testing.T) {
 				},
 			},
 		},
+
+		// Set, with list
+		{
+			Schema: map[string]*Schema{
+				"ports": &Schema{
+					Type:     TypeSet,
+					Optional: true,
+					Computed: true,
+					Elem:     &Schema{Type: TypeInt},
+					Set: func(a interface{}) int {
+						return a.(int)
+					},
+				},
+			},
+
+			State: &terraform.ResourceState{
+				Attributes: map[string]string{
+					"ports.#": "3",
+					"ports.0": "100",
+					"ports.1": "80",
+					"ports.2": "80",
+				},
+			},
+
+			Key:   "ports",
+			Value: []interface{}{100, 125, 125},
+
+			GetKey:   "ports",
+			GetValue: []interface{}{100, 125},
+		},
+
+		// Set, with Set
+		{
+			Schema: map[string]*Schema{
+				"ports": &Schema{
+					Type:     TypeSet,
+					Optional: true,
+					Computed: true,
+					Elem:     &Schema{Type: TypeInt},
+					Set: func(a interface{}) int {
+						return a.(int)
+					},
+				},
+			},
+
+			State: &terraform.ResourceState{
+				Attributes: map[string]string{
+					"ports.#": "3",
+					"ports.0": "100",
+					"ports.1": "80",
+					"ports.2": "80",
+				},
+			},
+
+			Key: "ports",
+			Value: &Set{
+				m: map[int]interface{}{
+					1: 1,
+					2: 2,
+				},
+			},
+
+			GetKey:   "ports",
+			GetValue: []interface{}{1, 2},
+		},
+
+		// Set single item
+		{
+			Schema: map[string]*Schema{
+				"ports": &Schema{
+					Type:     TypeSet,
+					Optional: true,
+					Computed: true,
+					Elem:     &Schema{Type: TypeInt},
+					Set: func(a interface{}) int {
+						return a.(int)
+					},
+				},
+			},
+
+			State: &terraform.ResourceState{
+				Attributes: map[string]string{
+					"ports.#": "2",
+					"ports.0": "100",
+					"ports.1": "80",
+				},
+			},
+
+			Key:   "ports.0",
+			Value: 256,
+			Err:   true,
+
+			GetKey:   "ports",
+			GetValue: []interface{}{80, 100},
+		},
 	}
 
 	for i, tc := range cases {
@@ -1069,6 +1164,9 @@ func TestResourceDataSet(t *testing.T) {
 		}
 
 		v := d.Get(tc.GetKey)
+		if s, ok := v.(*Set); ok {
+			v = s.List()
+		}
 		if !reflect.DeepEqual(v, tc.GetValue) {
 			t.Fatalf("Get Bad: %d\n\n%#v", i, v)
 		}