From 55818e64bc62143aac7ad6dbb00ba776740e7bbd Mon Sep 17 00:00:00 2001 From: Andrew Starr-Bochicchio Date: Fri, 25 Oct 2024 12:56:57 -0400 Subject: [PATCH] droplets: support listing GPU Droplets. --- args.go | 3 ++ commands/droplets.go | 14 ++++++++- commands/droplets_test.go | 19 ++++++++++++ do/droplets.go | 30 +++++++++++++------ do/mocks/DropletsService.go | 16 ++++++++++ integration/droplet_list_test.go | 50 ++++++++++++++++++++++++++++++++ 6 files changed, 122 insertions(+), 10 deletions(-) diff --git a/args.go b/args.go index 416e3f1c3..4b5290bb1 100644 --- a/args.go +++ b/args.go @@ -565,4 +565,7 @@ const ( // ArgTokenValidationServer is the server used to validate an OAuth token ArgTokenValidationServer = "token-validation-server" + + // ArgGPUs specifies to list GPU Droplets + ArgGPUs = "gpus" ) diff --git a/commands/droplets.go b/commands/droplets.go index 12a93b305..139eda2ff 100644 --- a/commands/droplets.go +++ b/commands/droplets.go @@ -113,6 +113,7 @@ If you do not specify a region, the Droplet is created in the default region for aliasOpt("ls"), displayerType(&displayers.Droplet{})) AddStringFlag(cmdRunDropletList, doctl.ArgRegionSlug, "", "", "Retrieves a list of Droplets in a specified region") AddStringFlag(cmdRunDropletList, doctl.ArgTagName, "", "", "Retrieves a list of Droplets with the specified tag name") + AddBoolFlag(cmdRunDropletList, doctl.ArgGPUs, "", false, "List GPU Droplets only. By default, only non-GPU Droplets are returned.") cmdRunDropletList.Example = `The following example retrieves a list of all Droplets in the ` + "`" + `nyc1` + "`" + ` region: doctl compute droplet list --region nyc1` cmdDropletNeighbors := CmdBuilder(cmd, RunDropletNeighbors, "neighbors ", "List a Droplet's neighbors on your account", `Lists your Droplets that are on the same physical hardware, including the following details:`+dropletDetails, Writer, @@ -656,6 +657,15 @@ func RunDropletList(c *CmdConfig) error { return err } + gpus, err := c.Doit.GetBool(c.NS, doctl.ArgGPUs) + if err != nil { + return err + } + + if gpus && tagName != "" { + return fmt.Errorf("The --gpus and --tag-name flags are mutually exclusive.") + } + matches := make([]glob.Glob, 0, len(c.Args)) for _, globStr := range c.Args { g, err := glob.Compile(globStr) @@ -669,7 +679,9 @@ func RunDropletList(c *CmdConfig) error { var matchedList do.Droplets var list do.Droplets - if tagName == "" { + if gpus { + list, err = ds.ListWithGPUs() + } else if tagName == "" { list, err = ds.List() } else { list, err = ds.ListByTag(tagName) diff --git a/commands/droplets_test.go b/commands/droplets_test.go index b62e859fa..2cc7ceeda 100644 --- a/commands/droplets_test.go +++ b/commands/droplets_test.go @@ -398,6 +398,16 @@ func TestDropletsListByTag(t *testing.T) { }) } +func TestDropletsListGPUs(t *testing.T) { + withTestClient(t, func(config *CmdConfig, tm *tcMocks) { + tm.droplets.EXPECT().ListWithGPUs().Return(testDropletList, nil) + + config.Doit.Set(config.NS, doctl.ArgGPUs, true) + err := RunDropletList(config) + assert.NoError(t, err) + }) +} + func TestDropletsTag(t *testing.T) { withTestClient(t, func(config *CmdConfig, tm *tcMocks) { trr := &godo.TagResourcesRequest{ @@ -415,6 +425,15 @@ func TestDropletsTag(t *testing.T) { }) } +func TestDropletsListGPUsAndTagsExclusive(t *testing.T) { + withTestClient(t, func(config *CmdConfig, tm *tcMocks) { + config.Doit.Set(config.NS, doctl.ArgGPUs, true) + config.Doit.Set(config.NS, doctl.ArgTagName, "my-tag") + err := RunDropletList(config) + assert.Error(t, err) + }) +} + func TestDropletsTagMultiple(t *testing.T) { withTestClient(t, func(config *CmdConfig, tm *tcMocks) { trr := &godo.TagResourcesRequest{ diff --git a/do/droplets.go b/do/droplets.go index b812977a8..e3e325a15 100644 --- a/do/droplets.go +++ b/do/droplets.go @@ -53,6 +53,7 @@ type Kernels []Kernel type DropletsService interface { List() (Droplets, error) ListByTag(string) (Droplets, error) + ListWithGPUs() (Droplets, error) Get(int) (*Droplet, error) Create(*godo.DropletCreateRequest, bool) (*Droplet, error) CreateMultiple(*godo.DropletMultiCreateRequest) (Droplets, error) @@ -93,18 +94,25 @@ func (ds *dropletsService) List() (Droplets, error) { return si, resp, err } - si, err := PaginateResp(f) - if err != nil { - return nil, err - } + return ds.list(f) +} - list := make(Droplets, len(si)) - for i := range si { - a := si[i].(godo.Droplet) - list[i] = Droplet{Droplet: &a} +func (ds *dropletsService) ListWithGPUs() (Droplets, error) { + f := func(opt *godo.ListOptions) ([]any, *godo.Response, error) { + list, resp, err := ds.client.Droplets.ListWithGPUs(context.TODO(), opt) + if err != nil { + return nil, nil, err + } + + si := make([]any, len(list)) + for i := range list { + si[i] = list[i] + } + + return si, resp, err } - return list, nil + return ds.list(f) } func (ds *dropletsService) ListByTag(tagName string) (Droplets, error) { @@ -122,6 +130,10 @@ func (ds *dropletsService) ListByTag(tagName string) (Droplets, error) { return si, resp, err } + return ds.list(f) +} + +func (ds *dropletsService) list(f Generator) (Droplets, error) { si, err := PaginateResp(f) if err != nil { return nil, err diff --git a/do/mocks/DropletsService.go b/do/mocks/DropletsService.go index 8b36751fc..3f8162f1f 100644 --- a/do/mocks/DropletsService.go +++ b/do/mocks/DropletsService.go @@ -21,6 +21,7 @@ import ( type MockDropletsService struct { ctrl *gomock.Controller recorder *MockDropletsServiceMockRecorder + isgomock struct{} } // MockDropletsServiceMockRecorder is the mock recorder for MockDropletsService. @@ -188,6 +189,21 @@ func (mr *MockDropletsServiceMockRecorder) ListByTag(arg0 any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListByTag", reflect.TypeOf((*MockDropletsService)(nil).ListByTag), arg0) } +// ListWithGPUs mocks base method. +func (m *MockDropletsService) ListWithGPUs() (do.Droplets, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListWithGPUs") + ret0, _ := ret[0].(do.Droplets) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListWithGPUs indicates an expected call of ListWithGPUs. +func (mr *MockDropletsServiceMockRecorder) ListWithGPUs() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListWithGPUs", reflect.TypeOf((*MockDropletsService)(nil).ListWithGPUs)) +} + // Neighbors mocks base method. func (m *MockDropletsService) Neighbors(arg0 int) (do.Droplets, error) { m.ctrl.T.Helper() diff --git a/integration/droplet_list_test.go b/integration/droplet_list_test.go index 29d7e2244..0f89cb1a3 100644 --- a/integration/droplet_list_test.go +++ b/integration/droplet_list_test.go @@ -38,6 +38,14 @@ var _ = suite("compute/droplet/list", func(t *testing.T, when spec.G, it spec.S) q := req.URL.Query() tag := q.Get("tag_name") + dtype := q.Get("type") + + if tag != "" && dtype != "" { + w.Write([]byte(`{"id": "unprocessible_entity", "message":"mutually exclusive query parameters"}`)) + w.WriteHeader(http.StatusUnprocessableEntity) + return + } + if tag == "some-tag" { w.Write([]byte(`{}`)) return @@ -48,6 +56,10 @@ var _ = suite("compute/droplet/list", func(t *testing.T, when spec.G, it spec.S) return } + if dtype == "gpus" { + w.Write([]byte(dropletListGPUsResponse)) + } + w.Write([]byte(dropletListResponse)) default: dump, err := httputil.DumpRequest(req, true) @@ -76,6 +88,23 @@ var _ = suite("compute/droplet/list", func(t *testing.T, when spec.G, it spec.S) }) }) + when("the gpu flag is passed", func() { + it("lists gpu droplets", func() { + cmd := exec.Command(builtBinaryPath, + "-t", "some-magic-token", + "-u", server.URL, + "compute", + "droplet", + "list", + "--gpus", + ) + + output, err := cmd.CombinedOutput() + expect.NoError(err, fmt.Sprintf("received error output: %s", output)) + expect.Equal(strings.TrimSpace(dropletGPUListOutput), strings.TrimSpace(string(output))) + }) + }) + when("a region is provided", func() { it("filters the returned droplets by region", func() { cmd := exec.Command(builtBinaryPath, @@ -163,9 +192,30 @@ const ( }] }` + dropletListGPUsResponse = ` +{ + "droplets": [{ + "id": 1111, + "name": "gpu-droplet", + "image": { + "distribution": "Ubuntu", + "name": "gpu-h100x8-640gb-200" + }, + "region": { + "slug": "tor1" + }, + "status": "active" + }] +}` + dropletListOutput = ` ID Name Public IPv4 Private IPv4 Public IPv6 Memory VCPUs Disk Region Image VPC UUID Status Tags Features Volumes 1111 some-droplet-name 0 0 0 some-region-slug some-distro some-image-name active test,yes remotes some-volume-id +` + + dropletGPUListOutput = ` +ID Name Public IPv4 Private IPv4 Public IPv6 Memory VCPUs Disk Region Image VPC UUID Status Tags Features Volumes +1111 gpu-droplet 0 0 0 tor1 Ubuntu gpu-h100x8-640gb-200 active ` dropletListRegionOutput = `