Merge pull request digitalocean#31 from digitalocean/refactor/create-…

…types Add image/ssh key structs for explicit type on droplet create.
amccarthy1 · Mar 21, 2015 · ca5927e · ca5927e
2 parents 2973bd3 + f1ccbcd
commit ca5927e
Show file tree

Hide file tree

Showing 4 changed files with 67 additions and 18 deletions.
diff --git a/README.md b/README.md
@@ -49,7 +49,9 @@ createRequest := &godo.DropletCreateRequest{
     Name:   dropletName,
     Region: "nyc3",
     Size:   "512mb",
-    Image:  "ubuntu-14-04-x64",
+    Image: DropletCreateImage{
+        Slug: "ubuntu-14-04-x64",
+    },
 }
 
 newDroplet, _, err := client.Droplets.Create(createRequest)

diff --git a/droplets.go b/droplets.go
@@ -1,6 +1,9 @@
 package godo
 
-import "fmt"
+import (
+	"encoding/json"
+	"fmt"
+)
 
 const dropletBasePath = "v2/droplets"
 
@@ -56,17 +59,45 @@ type dropletsRoot struct {
 	Links    *Links    `json:"links"`
 }
 
+// DropletCreateImage identifies an image for the create request. It prefers slug over ID.
+type DropletCreateImage struct {
+	ID   int
+	Slug string
+}
+
+func (d DropletCreateImage) MarshalJSON() ([]byte, error) {
+	if d.Slug != "" {
+		return json.Marshal(d.Slug)
+	} else {
+		return json.Marshal(d.ID)
+	}
+}
+
+// DropletCreateSSHKey identifies a SSH Key for the create request. It prefers fingerprint over ID.
+type DropletCreateSSHKey struct {
+	ID          int
+	Fingerprint string
+}
+
+func (d DropletCreateSSHKey) MarshalJSON() ([]byte, error) {
+	if d.Fingerprint != "" {
+		return json.Marshal(d.Fingerprint)
+	} else {
+		return json.Marshal(d.ID)
+	}
+}
+
 // DropletCreateRequest represents a request to create a droplet.
 type DropletCreateRequest struct {
-	Name              string        `json:"name"`
-	Region            string        `json:"region"`
-	Size              string        `json:"size"`
-	Image             string        `json:"image"`
-	SSHKeys           []interface{} `json:"ssh_keys"`
-	Backups           bool          `json:"backups"`
-	IPv6              bool          `json:"ipv6"`
-	PrivateNetworking bool          `json:"private_networking"`
-	UserData          string        `json:"user_data"`
+	Name              string                `json:"name"`
+	Region            string                `json:"region"`
+	Size              string                `json:"size"`
+	Image             DropletCreateImage    `json:"image"`
+	SSHKeys           []DropletCreateSSHKey `json:"ssh_keys"`
+	Backups           bool                  `json:"backups"`
+	IPv6              bool                  `json:"ipv6"`
+	PrivateNetworking bool                  `json:"private_networking"`
+	UserData          string                `json:"user_data"`
 }
 
 func (d DropletCreateRequest) String() string {

diff --git a/droplets_test.go b/droplets_test.go
@@ -126,16 +126,32 @@ func TestDroplets_Create(t *testing.T) {
 		Name:   "name",
 		Region: "region",
 		Size:   "size",
-		Image:  "1",
+		Image: DropletCreateImage{
+			ID: 1,
+		},
 	}
 
 	mux.HandleFunc("/v2/droplets", func(w http.ResponseWriter, r *http.Request) {
-		v := new(DropletCreateRequest)
-		json.NewDecoder(r.Body).Decode(v)
+		expected := map[string]interface{}{
+			"name":               "name",
+			"region":             "region",
+			"size":               "size",
+			"image":              float64(1),
+			"ssh_keys":           nil,
+			"backups":            false,
+			"ipv6":               false,
+			"private_networking": false,
+			"user_data":          "",
+		}
+
+		var v map[string]interface{}
+		err := json.NewDecoder(r.Body).Decode(&v)
+		if err != nil {
+			t.Fatal(err)
+		}
 
-		testMethod(t, r, "POST")
-		if !reflect.DeepEqual(v, createRequest) {
-			t.Errorf("Request body = %+v, expected %+v", v, createRequest)
+		if !reflect.DeepEqual(v, expected) {
+			t.Errorf("Request body = %#v, expected %#v", v, expected)
 		}
 
 		fmt.Fprintf(w, `{"droplet":{"id":1}, "links":{"actions": [{"id": 1, "href": "http://example.com", "rel": "create"}]}}`)

diff --git a/godo_test.go b/godo_test.go
@@ -79,7 +79,7 @@ func TestNewRequest(t *testing.T) {
 
 	inURL, outURL := "/foo", defaultBaseURL+"foo"
 	inBody, outBody := &DropletCreateRequest{Name: "l"},
-		`{"name":"l","region":"","size":"","image":"",`+
+		`{"name":"l","region":"","size":"","image":0,`+
 			`"ssh_keys":null,"backups":false,"ipv6":false,`+
 			`"private_networking":false,"user_data":""}`+"\n"
 	req, _ := c.NewRequest("GET", inURL, inBody)