diff --git a/acceptance/experimental/air/cancel/out.test.toml b/acceptance/experimental/air/cancel/out.test.toml new file mode 100644 index 0000000000..d6187dcb04 --- /dev/null +++ b/acceptance/experimental/air/cancel/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/cancel/output.txt b/acceptance/experimental/air/cancel/output.txt new file mode 100644 index 0000000000..9fd8a055f1 --- /dev/null +++ b/acceptance/experimental/air/cancel/output.txt @@ -0,0 +1,29 @@ + +=== cancel by id (text) +>>> [CLI] experimental air cancel 123 +Successfully requested cancellation for run 123 + +=== cancel by id (json) +>>> [CLI] experimental air cancel 123 -o json +{ + "v": 1, + "ts": "[TIMESTAMP]", + "data": { + "cancelled": [ + "123" + ] + } +} + +=== cancel multiple ids +>>> [CLI] experimental air cancel 123 456 +Successfully requested cancellation for run 123 +Successfully requested cancellation for run 456 +Successfully requested cancellation for 2 run(s). + +=== cancel --all +>>> [CLI] experimental air cancel --all -y +Searching active runs for [USERNAME] in [DATABRICKS_URL]... +Successfully requested cancellation for run [NUMID] +Successfully requested cancellation for run [NUMID] +Successfully requested cancellation for 2 run(s). diff --git a/acceptance/experimental/air/cancel/script b/acceptance/experimental/air/cancel/script new file mode 100644 index 0000000000..ce04a8977f --- /dev/null +++ b/acceptance/experimental/air/cancel/script @@ -0,0 +1,11 @@ +title "cancel by id (text)" +trace $CLI experimental air cancel 123 + +title "cancel by id (json)" +trace $CLI experimental air cancel 123 -o json + +title "cancel multiple ids" +trace $CLI experimental air cancel 123 456 + +title "cancel --all" +trace $CLI experimental air cancel --all -y diff --git a/acceptance/experimental/air/cancel/test.toml b/acceptance/experimental/air/cancel/test.toml new file mode 100644 index 0000000000..c73594501d --- /dev/null +++ b/acceptance/experimental/air/cancel/test.toml @@ -0,0 +1,53 @@ +# This command does not deploy a bundle, so no engine matrix is needed. +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# The SDK occasionally probes host reachability with a HEAD request; stub it so +# the test is deterministic. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +# CancelRun accepts the request and returns an empty body. +[[Server]] +Pattern = "POST /api/2.2/jobs/runs/cancel" +Response.Body = '{}' + +# Jobs runs/list backs `cancel --all`: two active AIR runs for the current user +# (tester@databricks.com, from the built-in scim/v2/Me handler). +[[Server]] +Pattern = "GET /api/2.2/jobs/runs/list" +Response.Body = ''' +{ + "runs": [ + { + "run_id": 334747067049496, + "run_name": "qwen-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717608759000, + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{ + "run_id": 334747067049497, + "ai_runtime_task": { + "experiment": "/Users/tester@databricks.com/qwen-train", + "deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}}] + } + }] + }, + { + "run_id": 566001814929041, + "run_name": "llama-train", + "creator_user_name": "tester@databricks.com", + "start_time": 1717612404000, + "state": {"life_cycle_state": "RUNNING"}, + "tasks": [{ + "run_id": 566001814929042, + "ai_runtime_task": { + "experiment": "/Users/tester@databricks.com/llama-train", + "deployments": [{"compute": {"accelerator_type": "GPU_1xA10", "accelerator_count": 1}}] + } + }] + } + ] +} +''' diff --git a/acceptance/experimental/air/unimplemented/output.txt b/acceptance/experimental/air/unimplemented/output.txt index 21c3c891af..7db6ef1aec 100644 --- a/acceptance/experimental/air/unimplemented/output.txt +++ b/acceptance/experimental/air/unimplemented/output.txt @@ -5,12 +5,6 @@ Error: `air logs` is not implemented yet Exit code: 1 -=== cancel ->>> [CLI] experimental air cancel 123 -Error: `air cancel` is not implemented yet - -Exit code: 1 - === register-image >>> [CLI] experimental air register-image my-image:latest Error: `air register-image` is not implemented yet diff --git a/acceptance/experimental/air/unimplemented/script b/acceptance/experimental/air/unimplemented/script index 4c53586b16..19dc13ffe8 100644 --- a/acceptance/experimental/air/unimplemented/script +++ b/acceptance/experimental/air/unimplemented/script @@ -3,8 +3,5 @@ title "logs" errcode trace $CLI experimental air logs 123 -title "cancel" -errcode trace $CLI experimental air cancel 123 - title "register-image" errcode trace $CLI experimental air register-image my-image:latest diff --git a/experimental/air/cmd/cancel.go b/experimental/air/cmd/cancel.go index ae5514e5b0..519a7a8206 100644 --- a/experimental/air/cmd/cancel.go +++ b/experimental/air/cmd/cancel.go @@ -1,10 +1,40 @@ package aircmd import ( + "context" + "errors" + "fmt" + "net/http" + "strconv" + "strings" + "text/tabwriter" + "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/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/spf13/cobra" ) +// cancelData is the JSON payload printed by `air cancel`. `all` is set only for +// --all, `workspace` only when --all finds no active runs, and `failed` only +// when a run could not be cancelled. +type cancelData struct { + Cancelled []string `json:"cancelled"` + All bool `json:"all,omitempty"` + Workspace string `json:"workspace,omitempty"` + Failed []cancelFailure `json:"failed,omitempty"` +} + +type cancelFailure struct { + RunID string `json:"run_id"` + Error string `json:"error"` +} + func newCancelCommand() *cobra.Command { var ( all bool @@ -15,9 +45,6 @@ func newCancelCommand() *cobra.Command { Use: "cancel [JOB_RUN_ID...]", Short: "Cancel one or more runs", Long: `Cancel one or more runs by ID, or cancel all of your active runs with --all.`, - RunE: func(cmd *cobra.Command, args []string) error { - return notImplemented("cancel") - }, } cmd.Flags().BoolVar(&all, "all", false, "Cancel all of your active runs") @@ -35,5 +62,161 @@ func newCancelCommand() *cobra.Command { return nil } + // In -o json mode an auth failure should be a JSON error envelope, not a bare + // error. ErrAlreadyPrinted passes through (already handled upstream). + 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 renderError(cmd.Context(), cmd, "INTERNAL_ERROR", "TRANSIENT", true, err) + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + w := cmdctx.WorkspaceClient(ctx) + jsonOut := root.OutputType(cmd) == flags.OutputJSON + + runIDs := args + data := cancelData{Cancelled: []string{}} + + if all { + data.All = true + + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to resolve current user: %w", err)) + } + host := strings.TrimRight(w.Config.Host, "/") + + if !jsonOut { + cmdio.LogString(ctx, fmt.Sprintf("Searching active runs for %s in %s...", me.UserName, host)) + } + + // Fetch every active run (up to the scan bound) so --all cancels all + // of them, not just the first page. + fetcher := newRunFetcher(ctx, w, listQuery{activeOnly: true, userFilter: me.UserName}) + rows, err := fetcher.next(maxListScan) + if err != nil { + return renderError(ctx, cmd, "INTERNAL_ERROR", "TRANSIENT", true, + fmt.Errorf("failed to list active runs: %w", err)) + } + + runIDs = make([]string, 0, len(rows)) + for i := range rows { + if rows[i].RunID != "" { + runIDs = append(runIDs, rows[i].RunID) + } + } + + if len(runIDs) == 0 { + if jsonOut { + data.Workspace = host + return renderEnvelope(ctx, data) + } + cmdio.LogString(ctx, "No active runs found.") + return nil + } + + if !yes { + displayCancelPreview(ctx, rows, host) + confirmed, err := cmdio.AskYesOrNo(ctx, fmt.Sprintf("\nCancel %d run(s) in %s?", len(runIDs), host)) + if err != nil { + return err + } + if !confirmed { + cmdio.LogString(ctx, "Cancellation aborted.") + return root.ErrAlreadyPrinted + } + } + } + + for _, rid := range runIDs { + err := cancelRun(ctx, w, rid) + if err != nil { + data.Failed = append(data.Failed, cancelFailure{RunID: rid, Error: err.Error()}) + if !jsonOut { + if runNotFound(err) { + cmdio.LogString(ctx, fmt.Sprintf("Run %s not found. Please check the run ID and ensure you're using a Job Run ID.", rid)) + } else { + cmdio.LogString(ctx, fmt.Sprintf("Failed to cancel run %s: %s", rid, err)) + } + } + continue + } + data.Cancelled = append(data.Cancelled, rid) + if !jsonOut { + cmdio.LogString(ctx, "Successfully requested cancellation for run "+rid) + } + } + + if jsonOut { + if err := renderEnvelope(ctx, data); err != nil { + return err + } + // Print the envelope, but still exit non-zero on any failure. + if len(data.Failed) > 0 { + return root.ErrAlreadyPrinted + } + return nil + } + + if len(data.Failed) > 0 { + cmdio.LogString(ctx, fmt.Sprintf("%d run(s) failed to cancel.", len(data.Failed))) + return root.ErrAlreadyPrinted + } + if all || len(data.Cancelled) > 1 { + cmdio.LogString(ctx, fmt.Sprintf("Successfully requested cancellation for %d run(s).", len(data.Cancelled))) + } + return nil + } + return cmd } + +// runNotFound reports whether err means the run does not exist. The cancel +// endpoint returns 400 INVALID_PARAMETER_VALUE ("Run does not exist") for +// an unknown run, and the SDK only remaps that to ErrResourceDoesNotExist for +// the runs/get path, not cancel — so we also detect the raw code here. +func runNotFound(err error) bool { + if errors.Is(err, apierr.ErrResourceDoesNotExist) { + return true + } + if apiErr, ok := errors.AsType[*apierr.APIError](err); ok { + return apiErr.StatusCode == http.StatusBadRequest && apiErr.ErrorCode == "INVALID_PARAMETER_VALUE" + } + return false +} + +// cancelRun requests cancellation of a single job run. The cancel is async, so +// the returned waiter is ignored. +func cancelRun(ctx context.Context, w *databricks.WorkspaceClient, rid string) error { + runID, err := strconv.ParseInt(rid, 10, 64) + if err != nil || runID <= 0 { + return fmt.Errorf("invalid run ID %q: must be a positive integer", rid) + } + _, err = w.Jobs.CancelRun(ctx, jobs.CancelRun{RunId: runID}) + return err +} + +// displayCancelPreview shows the runs that `cancel --all` is about to terminate. +func displayCancelPreview(ctx context.Context, rows []listRow, host string) { + var sb strings.Builder + fmt.Fprintf(&sb, "\nWorkspace: %s\n", host) + fmt.Fprintf(&sb, "Found %d active run(s) to cancel:\n\n", len(rows)) + + tw := tabwriter.NewWriter(&sb, 0, 0, 2, ' ', 0) + fmt.Fprintln(tw, "Run ID\tExperiment\tStarted") + for i := range rows { + experiment := orNA(rows[i].Experiment) + started := na + if rows[i].StartedAt != nil { + started = *rows[i].StartedAt + } + fmt.Fprintf(tw, "%s\t%s\t%s\n", rows[i].RunID, experiment, started) + } + tw.Flush() + + cmdio.LogString(ctx, strings.TrimRight(sb.String(), "\n")) +} diff --git a/experimental/air/cmd/cancel_test.go b/experimental/air/cmd/cancel_test.go new file mode 100644 index 0000000000..c676567003 --- /dev/null +++ b/experimental/air/cmd/cancel_test.go @@ -0,0 +1,289 @@ +package aircmd + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "testing/iotest" + + "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/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" +) + +// runCancelAll runs `cancel --all` against w with the given output mode and +// stdin, capturing output into buf. +func runCancelAll(t *testing.T, w *databricks.WorkspaceClient, out flags.Output, in io.Reader, buf *bytes.Buffer) error { + t.Helper() + cmd := withOutput(newCancelCommand(), out) + require.NoError(t, cmd.Flags().Set("all", "true")) + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, in, buf, buf, "", "")) + cmd.SetContext(cmdctx.SetWorkspaceClient(ctx, w)) + return cmd.RunE(cmd, nil) +} + +// cancelEnvelope decodes the air JSON envelope with the cancel payload. +type cancelEnvelope struct { + V int `json:"v"` + Data cancelData `json:"data"` +} + +// runCancel runs the cancel command against w with the given output mode and +// stdin, capturing stdout/stderr into buf. +func runCancel(t *testing.T, w *databricks.WorkspaceClient, out flags.Output, in string, buf *bytes.Buffer, args ...string) (*cobra.Command, error) { + t.Helper() + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), out, strings.NewReader(in), buf, buf, "", "")) + ctx = cmdctx.SetWorkspaceClient(ctx, w) + cmd := withOutput(newCancelCommand(), out) + cmd.SetContext(ctx) + return cmd, cmd.RunE(cmd, args) +} + +func TestCancelArgs(t *testing.T) { + tests := []struct { + name string + all bool + args []string + wantErr string + }{ + {name: "one id", args: []string{"123"}}, + {name: "many ids", args: []string{"123", "456"}}, + {name: "all", all: true}, + {name: "no input", wantErr: "provide at least one JOB_RUN_ID, or use --all"}, + {name: "ids with all", all: true, args: []string{"123"}, wantErr: "cannot combine JOB_RUN_ID arguments with --all"}, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + cmd := newCancelCommand() + if tc.all { + require.NoError(t, cmd.Flags().Set("all", "true")) + } + err := cmd.Args(cmd, tc.args) + if tc.wantErr == "" { + assert.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantErr) + }) + } +} + +func TestCancelRunInvalidID(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + for _, id := range []string{"abc", "0", "-1"} { + err := cancelRun(t.Context(), m.WorkspaceClient, id) + require.Error(t, err) + assert.Contains(t, err.Error(), "invalid run ID") + } +} + +func TestCancelByIDSuccess(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 123}).Return(nil, nil) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 456}).Return(nil, nil) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "123", "456") + require.NoError(t, err) + + out := buf.String() + assert.Contains(t, out, "Successfully requested cancellation for run 123") + assert.Contains(t, out, "Successfully requested cancellation for run 456") + // More than one run cancelled prints the count summary. + assert.Contains(t, out, "Successfully requested cancellation for 2 run(s).") +} + +func TestCancelByIDNotFound(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 5}).Return(nil, apierr.ErrResourceDoesNotExist) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "5") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + out := buf.String() + assert.Contains(t, out, "Run 5 not found") + assert.Contains(t, out, "1 run(s) failed to cancel.") +} + +func TestCancelByIDNotFoundInvalidParam(t *testing.T) { + // The cancel endpoint reports an unknown run as 400 INVALID_PARAMETER_VALUE, + // which the SDK does not remap to ErrResourceDoesNotExist for this path. + m := mocks.NewMockWorkspaceClient(t) + apiErr := &apierr.APIError{StatusCode: http.StatusBadRequest, ErrorCode: "INVALID_PARAMETER_VALUE", Message: "Run 5 does not exist."} + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 5}).Return(nil, apiErr) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "5") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "Run 5 not found") +} + +func TestCancelPartialFailureJSON(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 123}).Return(nil, nil) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 5}).Return(nil, apierr.ErrResourceDoesNotExist) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputJSON, "", &buf, "123", "5") + // The envelope is printed, but a failure still exits non-zero. + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + + var got cancelEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, []string{"123"}, got.Data.Cancelled) + require.Len(t, got.Data.Failed, 1) + assert.Equal(t, "5", got.Data.Failed[0].RunID) + assert.False(t, got.Data.All) +} + +func TestCancelByIDSuccessJSON(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 123}).Return(nil, nil) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputJSON, "", &buf, "123") + require.NoError(t, err) + + var got cancelEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Equal(t, []string{"123"}, got.Data.Cancelled) + assert.Empty(t, got.Data.Failed) +} + +func TestCancelByIDGenericFailure(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockJobsAPI().EXPECT().CancelRun(mock.Anything, jobs.CancelRun{RunId: 7}).Return(nil, errors.New("boom")) + + var buf bytes.Buffer + _, err := runCancel(t, m.WorkspaceClient, flags.OutputText, "", &buf, "7") + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "Failed to cancel run 7: boom") +} + +func TestCancelAllNoActiveRuns(t *testing.T) { + w := newTestWorkspaceClient(t, runsServer(t, runsListBody(t, "")).URL) + var buf bytes.Buffer + require.NoError(t, runCancelAll(t, w, flags.OutputText, nil, &buf)) + assert.Contains(t, buf.String(), "No active runs found.") +} + +func TestCancelAllNoActiveRunsJSON(t *testing.T) { + srv := runsServer(t, runsListBody(t, "")) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + require.NoError(t, runCancelAll(t, w, flags.OutputJSON, nil, &buf)) + + var got cancelEnvelope + require.NoError(t, json.Unmarshal(buf.Bytes(), &got)) + assert.Empty(t, got.Data.Cancelled) + assert.True(t, got.Data.All) + assert.Equal(t, srv.URL, got.Data.Workspace) +} + +func TestCancelAllConfirmYes(t *testing.T) { + srv := runsServer(t, runsListBody(t, "", + airJobRun(111, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-a"), + airJobRun(222, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-b"), + )) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + require.NoError(t, runCancelAll(t, w, flags.OutputText, strings.NewReader("y\n"), &buf)) + out := buf.String() + assert.Contains(t, out, "active run(s) to cancel") + assert.Contains(t, out, "Successfully requested cancellation for run 111") + assert.Contains(t, out, "Successfully requested cancellation for run 222") +} + +func TestCancelAllAbort(t *testing.T) { + srv := runsServer(t, runsListBody(t, "", + airJobRun(111, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-a"), + )) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + err := runCancelAll(t, w, flags.OutputText, strings.NewReader("n\n"), &buf) + require.ErrorIs(t, err, root.ErrAlreadyPrinted) + assert.Contains(t, buf.String(), "Cancellation aborted.") +} + +func TestCancelAllConfirmReadError(t *testing.T) { + srv := runsServer(t, runsListBody(t, "", + airJobRun(111, "me@example.com", "GPU_1xA10", 1, "/Users/me@example.com/exp-a"), + )) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + err := runCancelAll(t, w, flags.OutputText, iotest.ErrReader(errors.New("read failed")), &buf) + require.Error(t, err) + assert.Contains(t, err.Error(), "read failed") +} + +func TestCancelAllMeError(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, iam.MeRequest{}).Return(nil, errors.New("nope")) + + var buf bytes.Buffer + err := runCancelAll(t, m.WorkspaceClient, flags.OutputText, nil, &buf) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to resolve current user") +} + +func TestCancelAllListError(t *testing.T) { + // Me succeeds (default empty user), but listing active runs fails. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == jobsRunsListPath { + w.WriteHeader(http.StatusInternalServerError) + _, _ = w.Write([]byte(`{"error_code":"INTERNAL","message":"boom"}`)) + return + } + _, _ = w.Write([]byte(`{}`)) + })) + t.Cleanup(srv.Close) + w := newTestWorkspaceClient(t, srv.URL) + + var buf bytes.Buffer + err := runCancelAll(t, w, flags.OutputText, nil, &buf) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to list active runs") +} + +func TestDisplayCancelPreview(t *testing.T) { + var buf bytes.Buffer + ctx := cmdio.InContext(t.Context(), cmdio.NewIO(t.Context(), flags.OutputText, nil, &buf, &buf, "", "")) + + started := "2026-06-05 17:32 UTC" + rows := []listRow{ + {RunID: "111", Experiment: "exp-a", StartedAt: &started}, + {RunID: "222"}, // no experiment or start time -> N/A + } + displayCancelPreview(ctx, rows, "https://my-workspace.cloud.databricks.test") + + out := buf.String() + assert.Contains(t, out, "Workspace: https://my-workspace.cloud.databricks.test") + assert.Contains(t, out, "Found 2 active run(s) to cancel:") + assert.Contains(t, out, "Run ID") + assert.Contains(t, out, "111") + assert.Contains(t, out, "exp-a") + assert.Contains(t, out, "222") + assert.Contains(t, out, na) +} diff --git a/experimental/air/cmd/stubs_test.go b/experimental/air/cmd/stubs_test.go index 4607d7d9ea..e28d7f6673 100644 --- a/experimental/air/cmd/stubs_test.go +++ b/experimental/air/cmd/stubs_test.go @@ -14,7 +14,6 @@ import ( func TestStubCommandsReturnNotImplemented(t *testing.T) { stubs := map[string]*cobra.Command{ "logs": newLogsCommand(), - "cancel": newCancelCommand(), "register-image": newRegisterImageCommand(), }