From e99f22a846d05b3fa97bcec513310c6e985f297c Mon Sep 17 00:00:00 2001
From: Andrew Starr-Bochicchio <a.starr.b@gmail.com>
Date: Thu, 24 Oct 2024 18:41:31 -0400
Subject: [PATCH] droplets: support listing GPU Droplets.

---
 droplets.go      | 12 +++++++
 droplets_test.go | 92 ++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+)

diff --git a/droplets.go b/droplets.go
index 5f198636..1ed09ec8 100644
--- a/droplets.go
+++ b/droplets.go
@@ -17,6 +17,7 @@ var errNoNetworks = errors.New("no networks have been defined")
 // See: https://docs.digitalocean.com/reference/api/api-reference/#tag/Droplets
 type DropletsService interface {
 	List(context.Context, *ListOptions) ([]Droplet, *Response, error)
+	ListWithGPUs(context.Context, *ListOptions) ([]Droplet, *Response, error)
 	ListByName(context.Context, string, *ListOptions) ([]Droplet, *Response, error)
 	ListByTag(context.Context, string, *ListOptions) ([]Droplet, *Response, error)
 	Get(context.Context, int) (*Droplet, *Response, error)
@@ -321,6 +322,17 @@ func (s *DropletsServiceOp) List(ctx context.Context, opt *ListOptions) ([]Dropl
 	return s.list(ctx, path)
 }
 
+// ListWithGPUs lists all Droplets with GPUs.
+func (s *DropletsServiceOp) ListWithGPUs(ctx context.Context, opt *ListOptions) ([]Droplet, *Response, error) {
+	path := fmt.Sprintf("%s?type=gpus", dropletBasePath)
+	path, err := addOptions(path, opt)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return s.list(ctx, path)
+}
+
 // ListByName lists all Droplets filtered by name returning only exact matches.
 // It is case-insensitive
 func (s *DropletsServiceOp) ListByName(ctx context.Context, name string, opt *ListOptions) ([]Droplet, *Response, error) {
diff --git a/droplets_test.go b/droplets_test.go
index 17a62d75..ddd23f9e 100644
--- a/droplets_test.go
+++ b/droplets_test.go
@@ -44,6 +44,98 @@ func TestDroplets_ListDroplets(t *testing.T) {
 	}
 }
 
+func TestDroplets_ListDropletsWithGPUs(t *testing.T) {
+	setup()
+	defer teardown()
+
+	mux.HandleFunc("/v2/droplets", func(w http.ResponseWriter, r *http.Request) {
+		testMethod(t, r, http.MethodGet)
+		if r.URL.Query().Get("type") != "gpus" {
+			t.Errorf("Droplets.ListWithGPUs did not request with a type parameter")
+		}
+		fmt.Fprint(w, `{
+			"droplets": [
+				{
+					"id": 1,
+					"size": {
+						"gpu_info": {
+							"count": 1,
+							"vram": {
+								"amount": 8,
+								"unit": "gib"
+							},
+							"model": "nvidia_tesla_v100"
+						},
+						"disk_info": [
+							{
+								"type": "local",
+								"size": {
+									"amount": 200,
+									"unit": "gib"
+								}
+							},
+							{
+								"type": "scratch",
+								"size": {
+									"amount": 40960,
+									"unit": "gib"
+								}
+							}
+						]
+					}
+				}
+			],
+			"meta": {
+				"total": 1
+			}
+		}`)
+	})
+
+	droplets, resp, err := client.Droplets.ListWithGPUs(ctx, nil)
+	if err != nil {
+		t.Errorf("Droplets.List returned error: %v", err)
+	}
+
+	expectedDroplets := []Droplet{
+		{
+			ID: 1,
+			Size: &Size{
+				GPUInfo: &GPUInfo{
+					Count: 1,
+					VRAM: &VRAM{
+						Amount: 8,
+						Unit:   "gib",
+					},
+					Model: "nvidia_tesla_v100",
+				},
+				DiskInfo: []DiskInfo{
+					{
+						Type: "local",
+						Size: &DiskSize{
+							Amount: 200,
+							Unit:   "gib",
+						},
+					},
+					{
+						Type: "scratch",
+						Size: &DiskSize{
+							Amount: 40960,
+							Unit:   "gib",
+						},
+					},
+				},
+			},
+		},
+	}
+	if !reflect.DeepEqual(droplets, expectedDroplets) {
+		t.Errorf("Droplets.List\nDroplets: got=%#v\nwant=%#v", droplets, expectedDroplets)
+	}
+	expectedMeta := &Meta{Total: 1}
+	if !reflect.DeepEqual(resp.Meta, expectedMeta) {
+		t.Errorf("Droplets.List\nMeta: got=%#v\nwant=%#v", resp.Meta, expectedMeta)
+	}
+}
+
 func TestDroplets_ListDropletsByTag(t *testing.T) {
 	setup()
 	defer teardown()