diff --git a/acceptance/experimental/air/help/output.txt b/acceptance/experimental/air/help/output.txt index f2ab58d0a38..a512a1b0e8a 100644 --- a/acceptance/experimental/air/help/output.txt +++ b/acceptance/experimental/air/help/output.txt @@ -35,6 +35,10 @@ List your active runs for the current profile (use --all-status for finished run Usage: databricks experimental air list [flags] + databricks experimental air list [command] + +Available Commands: + provisioned-capacity List the pre-provisioned AI Runtime capacity reservations for the current workspace Flags: --all-status Show runs in all states (default: active only) @@ -49,6 +53,8 @@ Global Flags: -p, --profile string ~/.databrickscfg profile -t, --target string bundle target to use (if applicable) +Use "databricks experimental air list [command] --help" for more information about a command. + === logs help >>> [CLI] experimental air logs --help Stream logs from an active run, or fetch logs from a completed run. diff --git a/acceptance/experimental/air/provisioned-capacity/out.test.toml b/acceptance/experimental/air/provisioned-capacity/out.test.toml new file mode 100644 index 00000000000..0938e678987 --- /dev/null +++ b/acceptance/experimental/air/provisioned-capacity/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/experimental/air/provisioned-capacity/output.txt b/acceptance/experimental/air/provisioned-capacity/output.txt new file mode 100644 index 00000000000..4db6eb6bc35 --- /dev/null +++ b/acceptance/experimental/air/provisioned-capacity/output.txt @@ -0,0 +1,49 @@ + +=== list provisioned capacity (text) +>>> [CLI] experimental air list provisioned-capacity +ID ACCELERATOR RESERVED +cap-8xh100-alpha GPU_8xH100 64 +cap-1xh100-beta GPU_1xH100 8 + +=== list provisioned capacity (json) +>>> [CLI] experimental air list provisioned-capacity -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "provisioned_capacities": [ + { + "provisioned_capacity_id": "cap-8xh100-alpha", + "accelerator_type": "GPU_8xH100", + "reserved_accelerators": 64 + }, + { + "provisioned_capacity_id": "cap-1xh100-beta", + "accelerator_type": "GPU_1xH100", + "reserved_accelerators": 8 + } + ] + } +} + +=== get provisioned capacity (text) +>>> [CLI] experimental air get provisioned-capacity cap-8xh100-alpha +Provisioned Capacity ID: cap-8xh100-alpha +Accelerator Type: GPU_8xH100 +Reserved Accelerators: 64 +Used Accelerators: 40 +Idle Accelerators: 24 + +=== get provisioned capacity (json) +>>> [CLI] experimental air get provisioned-capacity cap-8xh100-alpha -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "provisioned_capacity_id": "cap-8xh100-alpha", + "accelerator_type": "GPU_8xH100", + "reserved_accelerators": 64, + "used_accelerators": 40, + "idle_accelerators": 24 + } +} diff --git a/acceptance/experimental/air/provisioned-capacity/script b/acceptance/experimental/air/provisioned-capacity/script new file mode 100644 index 00000000000..551cabbc69c --- /dev/null +++ b/acceptance/experimental/air/provisioned-capacity/script @@ -0,0 +1,11 @@ +title "list provisioned capacity (text)" +trace $CLI experimental air list provisioned-capacity + +title "list provisioned capacity (json)" +trace $CLI experimental air list provisioned-capacity -o json + +title "get provisioned capacity (text)" +trace $CLI experimental air get provisioned-capacity cap-8xh100-alpha + +title "get provisioned capacity (json)" +trace $CLI experimental air get provisioned-capacity cap-8xh100-alpha -o json diff --git a/acceptance/experimental/air/provisioned-capacity/test.toml b/acceptance/experimental/air/provisioned-capacity/test.toml new file mode 100644 index 00000000000..03c91e4c3e8 --- /dev/null +++ b/acceptance/experimental/air/provisioned-capacity/test.toml @@ -0,0 +1,29 @@ +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# List returns the workspace's reservations. usage is intentionally absent here: +# the real endpoint populates it only on Get. name is the AIP resource name. +[[Server]] +Pattern = "GET /api/2.0/ai-training/provisioned-capacities" +Response.Body = ''' +{ + "provisioned_capacities": [ + {"name": "provisioned-capacities/cap-8xh100-alpha", "spec": {"accelerator_type": "GPU_8xH100", "accelerator_count": 64}}, + {"name": "provisioned-capacities/cap-1xh100-beta", "spec": {"accelerator_type": "GPU_1xH100", "accelerator_count": 8}} + ] +} +''' + +# Get returns one reservation with its usage summary populated. +[[Server]] +Pattern = "GET /api/2.0/ai-training/provisioned-capacities/cap-8xh100-alpha" +Response.Body = ''' +{ + "name": "provisioned-capacities/cap-8xh100-alpha", + "spec": {"accelerator_type": "GPU_8xH100", "accelerator_count": 64}, + "status": {"usage": {"used_accelerator_count": 40, "idle_accelerator_count": 24}} +} +''' diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index 98dd9a656b3..83debb1150d 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -205,6 +205,8 @@ func newGetCommand() *cobra.Command { return nil } + cmd.AddCommand(newGetProvisionedCapacityCommand()) + return cmd } diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go index dfd9fbb38de..eef6071e765 100644 --- a/experimental/air/cmd/get_test.go +++ b/experimental/air/cmd/get_test.go @@ -34,12 +34,15 @@ func renderGet(t *testing.T, data getData) string { } // TestGetCommandShape locks in that `get` takes the run id directly as -// `air get JOB_RUN_ID` and has no `run` subcommand (it was collapsed back into -// `get`). The acceptance test exercises the happy path end to end. +// `air get JOB_RUN_ID` (there is no `run` subcommand — it was collapsed back +// into `get`). Its only subcommand is the `provisioned-capacity` noun. The +// acceptance test exercises the happy path end to end. func TestGetCommandShape(t *testing.T) { cmd := newGetCommand() assert.Equal(t, "get JOB_RUN_ID", cmd.Use) - assert.Empty(t, cmd.Commands(), "get must not register subcommands") + subs := cmd.Commands() + require.Len(t, subs, 1) + assert.Equal(t, "provisioned-capacity", subs[0].Name()) // ExactArgs(1): exactly one run id is required. assert.NoError(t, cmd.Args(cmd, []string{"123"})) assert.Error(t, cmd.Args(cmd, []string{})) diff --git a/experimental/air/cmd/list.go b/experimental/air/cmd/list.go index e0c3c3da483..59ffc840d5b 100644 --- a/experimental/air/cmd/list.go +++ b/experimental/air/cmd/list.go @@ -160,6 +160,8 @@ func newListCommand() *cobra.Command { return renderListText(cmd, fetcher, limit) } + cmd.AddCommand(newListProvisionedCapacityCommand()) + return cmd } diff --git a/experimental/air/cmd/provisioned_capacity.go b/experimental/air/cmd/provisioned_capacity.go new file mode 100644 index 00000000000..2eca0a6f4ed --- /dev/null +++ b/experimental/air/cmd/provisioned_capacity.go @@ -0,0 +1,312 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "strconv" + "strings" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/flags" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/client" + "github.com/spf13/cobra" +) + +// A "provisioned capacity" is a pre-provisioned AI Runtime accelerator +// reservation — the same reservation a run targets via +// compute.provisioned_capacity_id. It is served by AiWorkflowService's +// purpose-built public read API (the platform tracks it internally as a +// "guaranteed capacity"; the public surface calls it "provisioned capacity"). +const provisionedCapacityPath = "/api/2.0/ai-training/provisioned-capacities" + +// resourceNamePrefix is the AIP resource-name prefix on ProvisionedCapacity.name +// ("provisioned-capacities/{id}"). The CLI shows and accepts the bare id. +const resourceNamePrefix = "provisioned-capacities/" + +// capacityListPageSize is the per-request page size for the list endpoint. A +// workspace holds very few reservations, so this is effectively a single page. +const capacityListPageSize = 100 + +// capacityListMaxPages caps pagination so a misbehaving next_page_token can't +// loop forever. +const capacityListMaxPages = 50 + +// apiInt64 decodes an int64 that the REST gateway may send either as a JSON +// number or, per proto3 JSON, as a quoted string. +type apiInt64 int64 + +func (n *apiInt64) UnmarshalJSON(b []byte) error { + s := strings.Trim(string(b), `"`) + if s == "" || s == "null" { + return nil + } + v, err := strconv.ParseInt(s, 10, 64) + if err != nil { + return err + } + *n = apiInt64(v) + return nil +} + +// The wire structs below mirror the ProvisionedCapacity proto's JSON. +type provisionedCapacity struct { + Name string `json:"name"` + Spec *provisionedCapacitySpec `json:"spec"` + Status *provisionedCapacityStatus `json:"status"` +} + +type provisionedCapacitySpec struct { + AcceleratorType string `json:"accelerator_type"` + AcceleratorCount apiInt64 `json:"accelerator_count"` +} + +type provisionedCapacityStatus struct { + Usage *provisionedCapacityUsage `json:"usage"` +} + +type provisionedCapacityUsage struct { + UsedAcceleratorCount apiInt64 `json:"used_accelerator_count"` + IdleAcceleratorCount apiInt64 `json:"idle_accelerator_count"` +} + +type listProvisionedCapacitiesResponse struct { + ProvisionedCapacities []provisionedCapacity `json:"provisioned_capacities"` + NextPageToken string `json:"next_page_token"` +} + +// capacityListData is the `air list provisioned-capacity` payload. Usage counts +// are intentionally absent: the list endpoint does not populate them (they come +// from `air get`). +type capacityListData struct { + Rows []capacityRow `json:"provisioned_capacities"` +} + +type capacityRow struct { + ID string `json:"provisioned_capacity_id"` + AcceleratorType string `json:"accelerator_type"` + ReservedAccelerators int64 `json:"reserved_accelerators"` +} + +// capacityDetailData is the `air get provisioned-capacity` payload. Usage is a +// pointer because it is populated only when the reservation reports it. +type capacityDetailData struct { + ID string `json:"provisioned_capacity_id"` + AcceleratorType string `json:"accelerator_type"` + ReservedAccelerators int64 `json:"reserved_accelerators"` + UsedAccelerators *int64 `json:"used_accelerators"` + IdleAccelerators *int64 `json:"idle_accelerators"` +} + +// capacityID strips the AIP resource-name prefix, returning the bare id. Input +// may already be the bare id (from a user argument) or the full resource name +// (from a response's name field). +func capacityID(name string) string { + return strings.TrimPrefix(name, resourceNamePrefix) +} + +// capacityAPIError classifies a provisioned-capacities call failure into the +// CLI's error envelope, matching how `air get` classifies run lookups. +func capacityAPIError(ctx context.Context, cmd *cobra.Command, action string, err error) error { + // The handler gates the API behind a SAFE flag and reports FEATURE_DISABLED + // where it is not yet rolled out. That is a permanent state for the + // workspace, not a retryable failure — and its explicit error code is more + // specific than the generic 403 it may arrive as, so check it first. + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok && apiErr.ErrorCode == "FEATURE_DISABLED" { + return renderError(ctx, cmd, "FEATURE_DISABLED", "PERMANENT", false, + errors.New("the provisioned capacity API is not enabled for this workspace")) + } + if errors.Is(err, apierr.ErrUnauthenticated) || errors.Is(err, apierr.ErrPermissionDenied) { + return authError(ctx, cmd, err) + } + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to %s: %w", action, err)) +} + +// listProvisionedCapacities pages the list endpoint fully. +func listProvisionedCapacities(ctx context.Context, w *databricks.WorkspaceClient) ([]provisionedCapacity, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + + var out []provisionedCapacity + pageToken := "" + for range capacityListMaxPages { + // GET query params ride the request arg (the SDK serializes them for a + // GET), matching the sibling workflows call. + query := map[string]any{"page_size": capacityListPageSize} + if pageToken != "" { + query["page_token"] = pageToken + } + var resp listProvisionedCapacitiesResponse + if err := apiClient.Do(ctx, http.MethodGet, provisionedCapacityPath, nil, nil, query, &resp); err != nil { + return nil, err + } + out = append(out, resp.ProvisionedCapacities...) + if resp.NextPageToken == "" { + break + } + pageToken = resp.NextPageToken + } + return out, nil +} + +// getProvisionedCapacity fetches one reservation, including its usage summary. +func getProvisionedCapacity(ctx context.Context, w *databricks.WorkspaceClient, id string) (*provisionedCapacity, error) { + apiClient, err := client.New(w.Config) + if err != nil { + return nil, fmt.Errorf("failed to create API client: %w", err) + } + var pc provisionedCapacity + if err := apiClient.Do(ctx, http.MethodGet, provisionedCapacityPath+"/"+id, nil, nil, nil, &pc); err != nil { + return nil, err + } + return &pc, nil +} + +func newListProvisionedCapacityCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "provisioned-capacity", + Args: root.NoArgs, + Short: "List the pre-provisioned AI Runtime capacity reservations for the current workspace", + } + + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + err := root.MustWorkspaceClient(cmd, args) + if err == nil || errors.Is(err, root.ErrAlreadyPrinted) { + return err + } + return authError(cmd.Context(), cmd, err) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + + capacities, err := listProvisionedCapacities(ctx, w) + if err != nil { + return capacityAPIError(ctx, cmd, "list provisioned capacities", err) + } + + data := capacityListData{Rows: make([]capacityRow, 0, len(capacities))} + for _, c := range capacities { + data.Rows = append(data.Rows, capacityRowFrom(c)) + } + + if root.OutputType(cmd) != flags.OutputText { + return renderEnvelope(ctx, data) + } + renderCapacityTable(cmd.OutOrStdout(), data.Rows) + return nil + } + + return cmd +} + +func newGetProvisionedCapacityCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "provisioned-capacity PROVISIONED_CAPACITY_ID", + Args: root.ExactArgs(1), + Short: "Show a pre-provisioned AI Runtime capacity reservation, including its accelerator usage", + } + + cmd.PreRunE = func(cmd *cobra.Command, args []string) error { + err := root.MustWorkspaceClient(cmd, args) + if err == nil || errors.Is(err, root.ErrAlreadyPrinted) { + return err + } + return authError(cmd.Context(), cmd, err) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + + // Accept either the bare id or the full resource name. + id := capacityID(strings.TrimSpace(args[0])) + if id == "" { + return renderError(ctx, cmd, "INVALID_ARGS", "PERMANENT", false, + errors.New("provisioned_capacity_id cannot be empty")) + } + + pc, err := getProvisionedCapacity(ctx, w, id) + if err != nil { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return renderError(ctx, cmd, "NOT_FOUND", "NOT_FOUND", false, + fmt.Errorf("provisioned capacity %q not found: check the id with `air list provisioned-capacity`", id)) + } + return capacityAPIError(ctx, cmd, fmt.Sprintf("get provisioned capacity %q", id), err) + } + + data := capacityDetailFrom(*pc) + if root.OutputType(cmd) != flags.OutputText { + return renderEnvelope(ctx, data) + } + renderCapacityDetail(cmd.OutOrStdout(), data) + return nil + } + + return cmd +} + +// capacityRowFrom projects a wire ProvisionedCapacity to a list row. +func capacityRowFrom(c provisionedCapacity) capacityRow { + row := capacityRow{ID: capacityID(c.Name)} + if c.Spec != nil { + row.AcceleratorType = c.Spec.AcceleratorType + row.ReservedAccelerators = int64(c.Spec.AcceleratorCount) + } + return row +} + +// capacityDetailFrom projects a wire ProvisionedCapacity to the detail payload. +func capacityDetailFrom(c provisionedCapacity) capacityDetailData { + data := capacityDetailData{ID: capacityID(c.Name)} + if c.Spec != nil { + data.AcceleratorType = c.Spec.AcceleratorType + data.ReservedAccelerators = int64(c.Spec.AcceleratorCount) + } + if c.Status != nil && c.Status.Usage != nil { + used, idle := int64(c.Status.Usage.UsedAcceleratorCount), int64(c.Status.Usage.IdleAcceleratorCount) + data.UsedAccelerators = &used + data.IdleAccelerators = &idle + } + return data +} + +// renderCapacityTable prints the list as an aligned text table. +func renderCapacityTable(out io.Writer, rows []capacityRow) { + if len(rows) == 0 { + fmt.Fprintln(out, "No provisioned capacity reservations found.") + return + } + fmt.Fprintf(out, "%-40s %-14s %s\n", "ID", "ACCELERATOR", "RESERVED") + for _, r := range rows { + fmt.Fprintf(out, "%-40s %-14s %d\n", orNA(r.ID), orNA(r.AcceleratorType), r.ReservedAccelerators) + } +} + +// renderCapacityDetail prints the get view as aligned label/value lines. +func renderCapacityDetail(out io.Writer, d capacityDetailData) { + line := func(label, value string) { fmt.Fprintf(out, "%-24s %s\n", label+":", value) } + line("Provisioned Capacity ID", orNA(d.ID)) + line("Accelerator Type", orNA(d.AcceleratorType)) + line("Reserved Accelerators", strconv.FormatInt(d.ReservedAccelerators, 10)) + line("Used Accelerators", acceleratorCell(d.UsedAccelerators)) + line("Idle Accelerators", acceleratorCell(d.IdleAccelerators)) +} + +// acceleratorCell renders an optional count, showing N/A when the reservation +// did not report usage. +func acceleratorCell(v *int64) string { + if v == nil { + return na + } + return strconv.FormatInt(*v, 10) +} diff --git a/experimental/air/cmd/provisioned_capacity_test.go b/experimental/air/cmd/provisioned_capacity_test.go new file mode 100644 index 00000000000..86ece921748 --- /dev/null +++ b/experimental/air/cmd/provisioned_capacity_test.go @@ -0,0 +1,268 @@ +package aircmd + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/databricks/cli/cmd/root" + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/flags" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const capacityBasePath = "/api/2.0/ai-training/provisioned-capacities" + +func TestProvisionedCapacityCommandShape(t *testing.T) { + list := newListProvisionedCapacityCommand() + assert.Equal(t, "provisioned-capacity", list.Use) + assert.NoError(t, list.Args(list, []string{})) + assert.Error(t, list.Args(list, []string{"x"})) + + get := newGetProvisionedCapacityCommand() + assert.Equal(t, "provisioned-capacity PROVISIONED_CAPACITY_ID", get.Use) + assert.NoError(t, get.Args(get, []string{"cap-1"})) + assert.Error(t, get.Args(get, []string{})) + + // The subcommands are wired under `air list` and `air get`. + assert.True(t, hasSubcommand(newListCommand(), "provisioned-capacity")) + assert.True(t, hasSubcommand(newGetCommand(), "provisioned-capacity")) +} + +func hasSubcommand(parent *cobra.Command, name string) bool { + for _, c := range parent.Commands() { + if c.Name() == name { + return true + } + } + return false +} + +func TestAPIInt64Unmarshal(t *testing.T) { + var s struct { + N apiInt64 `json:"n"` + } + require.NoError(t, json.Unmarshal([]byte(`{"n": 42}`), &s)) + assert.EqualValues(t, 42, s.N) + + // proto3 JSON encodes int64 as a string; accept it too. + require.NoError(t, json.Unmarshal([]byte(`{"n": "64"}`), &s)) + assert.EqualValues(t, 64, s.N) + + // A null (or absent) field leaves the zero value in place: encoding/json + // treats UnmarshalJSON("null") as a no-op, so a fresh struct stays 0. + var fresh struct { + N apiInt64 `json:"n"` + } + require.NoError(t, json.Unmarshal([]byte(`{"n": null}`), &fresh)) + assert.EqualValues(t, 0, fresh.N) +} + +func TestCapacityID(t *testing.T) { + // The full resource name is reduced to the bare id; a bare id is unchanged. + assert.Equal(t, "cap-a", capacityID("provisioned-capacities/cap-a")) + assert.Equal(t, "cap-a", capacityID("cap-a")) +} + +// capacityServer serves the list and get endpoints from the given bodies. The +// list body is returned for the collection path; getByID maps a bare id to its +// body (missing ids return 404 RESOURCE_DOES_NOT_EXIST). +func capacityServer(t *testing.T, listPages []string, getByID map[string]string) *httptest.Server { + t.Helper() + call := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == capacityBasePath: + body := listPages[min(call, len(listPages)-1)] + call++ + _, _ = w.Write([]byte(body)) + case strings.HasPrefix(r.URL.Path, capacityBasePath+"/"): + id := strings.TrimPrefix(r.URL.Path, capacityBasePath+"/") + body, ok := getByID[id] + if !ok { + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error_code":"RESOURCE_DOES_NOT_EXIST","message":"not found"}`)) + return + } + _, _ = w.Write([]byte(body)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv +} + +func TestListProvisionedCapacityJSON(t *testing.T) { + // Two pages exercise the pagination loop; usage is absent on list, as the + // real endpoint leaves it. name arrives as the full resource name. + page1 := `{"provisioned_capacities":[{"name":"provisioned-capacities/cap-a","spec":{"accelerator_type":"GPU_8xH100","accelerator_count":64}}],"next_page_token":"tok2"}` + page2 := `{"provisioned_capacities":[{"name":"provisioned-capacities/cap-b","spec":{"accelerator_type":"GPU_1xH100","accelerator_count":"8"}}]}` + srv := capacityServer(t, []string{page1, page2}, nil) + + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(t.Context(), newTestWorkspaceClient(t, srv.URL)) + ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", "")) + cmd := withOutput(newListProvisionedCapacityCommand(), flags.OutputJSON) + cmd.SetContext(ctx) + + require.NoError(t, cmd.RunE(cmd, nil)) + + var got struct { + Data capacityListData `json:"data"` + } + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + require.Len(t, got.Data.Rows, 2) + assert.Equal(t, capacityRow{ID: "cap-a", AcceleratorType: "GPU_8xH100", ReservedAccelerators: 64}, got.Data.Rows[0]) + // The second row's accelerator_count arrived as a JSON string, and the id is + // stripped from the resource name. + assert.Equal(t, "cap-b", got.Data.Rows[1].ID) + assert.Equal(t, int64(8), got.Data.Rows[1].ReservedAccelerators) +} + +func TestListProvisionedCapacityText(t *testing.T) { + page := `{"provisioned_capacities":[{"name":"provisioned-capacities/cap-a","spec":{"accelerator_type":"GPU_8xH100","accelerator_count":64}}]}` + srv := capacityServer(t, []string{page}, nil) + + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(newListProvisionedCapacityCommand(), flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + require.NoError(t, cmd.RunE(cmd, nil)) + out := buf.String() + assert.Contains(t, out, "ACCELERATOR") + assert.Contains(t, out, "cap-a") + assert.Contains(t, out, "GPU_8xH100") + assert.Contains(t, out, "64") +} + +func TestListProvisionedCapacityTextEmpty(t *testing.T) { + srv := capacityServer(t, []string{`{"provisioned_capacities":[]}`}, nil) + + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(newListProvisionedCapacityCommand(), flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + require.NoError(t, cmd.RunE(cmd, nil)) + assert.Contains(t, buf.String(), "No provisioned capacity reservations found.") +} + +func TestGetProvisionedCapacityJSON(t *testing.T) { + body := `{"name":"provisioned-capacities/cap-a","spec":{"accelerator_type":"GPU_8xH100","accelerator_count":64},"status":{"usage":{"used_accelerator_count":40,"idle_accelerator_count":24}}}` + srv := capacityServer(t, nil, map[string]string{"cap-a": body}) + + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(t.Context(), newTestWorkspaceClient(t, srv.URL)) + ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", "")) + cmd := withOutput(newGetProvisionedCapacityCommand(), flags.OutputJSON) + cmd.SetContext(ctx) + + require.NoError(t, cmd.RunE(cmd, []string{"cap-a"})) + + var got struct { + Data capacityDetailData `json:"data"` + } + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, "cap-a", got.Data.ID) + assert.Equal(t, "GPU_8xH100", got.Data.AcceleratorType) + assert.Equal(t, int64(64), got.Data.ReservedAccelerators) + require.NotNil(t, got.Data.UsedAccelerators) + assert.Equal(t, int64(40), *got.Data.UsedAccelerators) + require.NotNil(t, got.Data.IdleAccelerators) + assert.Equal(t, int64(24), *got.Data.IdleAccelerators) +} + +func TestGetProvisionedCapacityAcceptsResourceName(t *testing.T) { + // A full resource name argument resolves to the same GET path as the bare id. + body := `{"name":"provisioned-capacities/cap-a","spec":{"accelerator_type":"GPU_8xH100","accelerator_count":64},"status":{"usage":{"used_accelerator_count":1,"idle_accelerator_count":63}}}` + srv := capacityServer(t, nil, map[string]string{"cap-a": body}) + + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(newGetProvisionedCapacityCommand(), flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&bytes.Buffer{}) + + require.NoError(t, cmd.RunE(cmd, []string{"provisioned-capacities/cap-a"})) +} + +func TestGetProvisionedCapacityText(t *testing.T) { + // A reservation whose usage block is absent shows N/A for the usage cells + // rather than a misleading zero. + body := `{"name":"provisioned-capacities/cap-a","spec":{"accelerator_type":"GPU_8xH100","accelerator_count":64}}` + srv := capacityServer(t, nil, map[string]string{"cap-a": body}) + + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, srv.URL)) + cmd := withOutput(newGetProvisionedCapacityCommand(), flags.OutputText) + cmd.SetContext(ctx) + cmd.SetOut(&buf) + + require.NoError(t, cmd.RunE(cmd, []string{"cap-a"})) + out := buf.String() + assert.Contains(t, out, "Provisioned Capacity ID: cap-a") + assert.Contains(t, out, "Reserved Accelerators: 64") + assert.Contains(t, out, "Used Accelerators: N/A") +} + +func TestGetProvisionedCapacityEmptyID(t *testing.T) { + ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), newTestWorkspaceClient(t, "https://example.com")) + cmd := withOutput(newGetProvisionedCapacityCommand(), flags.OutputText) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, []string{" "}) + require.Error(t, err) + assert.Contains(t, err.Error(), "cannot be empty") +} + +func TestListProvisionedCapacityFeatureDisabledJSON(t *testing.T) { + // The SAFE-gated handler reports FEATURE_DISABLED before rollout; it must be + // a permanent, non-retryable error, not a transient one. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code":"FEATURE_DISABLED","message":"The provisioned capacity API is not yet enabled."}`)) + })) + t.Cleanup(srv.Close) + + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(t.Context(), newTestWorkspaceClient(t, srv.URL)) + ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", "")) + cmd := withOutput(newListProvisionedCapacityCommand(), flags.OutputJSON) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, nil) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + var got errorEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, "FEATURE_DISABLED", got.Error.Code) + assert.False(t, got.Error.Retryable) + assert.Contains(t, got.Error.Message, "not enabled for this workspace") +} + +func TestGetProvisionedCapacityNotFoundJSON(t *testing.T) { + srv := capacityServer(t, nil, map[string]string{}) + + var buf bytes.Buffer + ctx := cmdctx.SetWorkspaceClient(t.Context(), newTestWorkspaceClient(t, srv.URL)) + ctx = cmdio.InContext(ctx, cmdio.NewIO(ctx, flags.OutputJSON, nil, &buf, &buf, "", "")) + cmd := withOutput(newGetProvisionedCapacityCommand(), flags.OutputJSON) + cmd.SetContext(ctx) + + err := cmd.RunE(cmd, []string{"missing"}) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + var got errorEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, "NOT_FOUND", got.Error.Code) + assert.Contains(t, got.Error.Message, `provisioned capacity "missing" not found`) +}