|
1 | 1 | package auth |
2 | 2 |
|
3 | 3 | import ( |
| 4 | + "encoding/json" |
4 | 5 | "errors" |
| 6 | + "net/http" |
| 7 | + "net/http/httptest" |
5 | 8 | "os" |
6 | 9 | "path/filepath" |
7 | 10 | "testing" |
8 | 11 |
|
9 | 12 | "github.com/databricks/cli/libs/auth/storage" |
10 | 13 | "github.com/databricks/cli/libs/cmdctx" |
11 | 14 | "github.com/databricks/cli/libs/databrickscfg/profile" |
| 15 | + "github.com/databricks/databricks-sdk-go/apierr" |
12 | 16 | "github.com/databricks/databricks-sdk-go/config" |
13 | 17 | "github.com/databricks/databricks-sdk-go/experimental/mocks" |
14 | 18 | "github.com/databricks/databricks-sdk-go/service/iam" |
@@ -397,6 +401,166 @@ func TestGetWorkspaceAuthStatus_U2M_PopulatesTokenStorage(t *testing.T) { |
397 | 401 | assert.Equal(t, "DATABRICKS_AUTH_STORAGE environment variable", status.TokenStorage.Source) |
398 | 402 | } |
399 | 403 |
|
| 404 | +// describeVerifyServer simulates the host that describe's secondary |
| 405 | +// verification calls hit: CurrentUser.Me and account Workspaces.List. |
| 406 | +// Statuses other than 200 return an error body with that status code; a zero |
| 407 | +// status marks an endpoint that must not be called at all. |
| 408 | +func describeVerifyServer(t *testing.T, accountID string, meStatus, listStatus int) *httptest.Server { |
| 409 | + t.Helper() |
| 410 | + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 411 | + respond := func(status int, okBody any) { |
| 412 | + if status == 0 { |
| 413 | + t.Errorf("unexpected request to %s", r.URL.Path) |
| 414 | + status = http.StatusNotFound |
| 415 | + } |
| 416 | + w.Header().Set("Content-Type", "application/json") |
| 417 | + if status != http.StatusOK { |
| 418 | + w.WriteHeader(status) |
| 419 | + okBody = map[string]any{"error_code": "TEST_ERROR", "message": "secondary check failed"} |
| 420 | + } |
| 421 | + require.NoError(t, json.NewEncoder(w).Encode(okBody)) |
| 422 | + } |
| 423 | + switch r.URL.Path { |
| 424 | + case "/api/2.0/preview/scim/v2/Me": |
| 425 | + respond(meStatus, map[string]any{"userName": "fallback-user"}) |
| 426 | + case "/api/2.0/accounts/" + accountID + "/workspaces": |
| 427 | + respond(listStatus, []map[string]any{}) |
| 428 | + default: |
| 429 | + w.WriteHeader(http.StatusNotFound) |
| 430 | + } |
| 431 | + })) |
| 432 | + t.Cleanup(server.Close) |
| 433 | + return server |
| 434 | +} |
| 435 | + |
| 436 | +// newDescribeWorkspaceCmd wires a mocked workspace client whose Me call fails |
| 437 | +// with meErr, with cfg as its resolved configuration. |
| 438 | +func newDescribeWorkspaceCmd(t *testing.T, cfg *config.Config, meErr error) *cobra.Command { |
| 439 | + t.Helper() |
| 440 | + m := mocks.NewMockWorkspaceClient(t) |
| 441 | + m.WorkspaceClient.Config = cfg |
| 442 | + m.GetMockCurrentUserAPI().EXPECT().Me(mock.Anything, mock.Anything).Return(nil, meErr) |
| 443 | + cmd := newHostProfileCmd(t) |
| 444 | + cmd.SetContext(cmdctx.SetWorkspaceClient(t.Context(), m.WorkspaceClient)) |
| 445 | + t.Setenv("DATABRICKS_CONFIG_FILE", filepath.Join(t.TempDir(), ".databrickscfg")) |
| 446 | + return cmd |
| 447 | +} |
| 448 | + |
| 449 | +// newDescribeAccountCmd wires a mocked account client whose Workspaces.List |
| 450 | +// call fails with listErr, with cfg as its resolved configuration. |
| 451 | +func newDescribeAccountCmd(t *testing.T, cfg *config.Config, listErr error) *cobra.Command { |
| 452 | + t.Helper() |
| 453 | + m := mocks.NewMockAccountClient(t) |
| 454 | + m.AccountClient.Config = cfg |
| 455 | + m.GetMockWorkspacesAPI().EXPECT().List(mock.Anything).Return(nil, listErr) |
| 456 | + cmd := newHostProfileCmd(t) |
| 457 | + cmd.SetContext(cmdctx.SetAccountClient(t.Context(), m.AccountClient)) |
| 458 | + t.Setenv("DATABRICKS_CONFIG_FILE", filepath.Join(t.TempDir(), ".databrickscfg")) |
| 459 | + return cmd |
| 460 | +} |
| 461 | + |
| 462 | +// resolveCfg returns a tryAuth that resolves cfg from attrs and reports the |
| 463 | +// given client type, mirroring what root.MustAnyClient leaves behind. |
| 464 | +func resolveCfg(t *testing.T, cfg *config.Config, attrs map[string]string, isAccount bool) tryAuth { |
| 465 | + t.Helper() |
| 466 | + return func(cmd *cobra.Command, args []string) (*config.Config, bool, error) { |
| 467 | + require.NoError(t, config.ConfigAttributes.ResolveFromStringMap(cfg, attrs)) |
| 468 | + return cfg, isAccount, nil |
| 469 | + } |
| 470 | +} |
| 471 | + |
| 472 | +// TestGetAuthStatusVerificationFallback covers the fallback when the primary |
| 473 | +// verification call (mocked client) fails: describe tries the other endpoint |
| 474 | +// over HTTP against describeVerifyServer, and only when both fail does it |
| 475 | +// report the primary error. Error rows leave wantUsername/wantAccountID empty |
| 476 | +// because errorAuthStatus never sets them. |
| 477 | +func TestGetAuthStatusVerificationFallback(t *testing.T) { |
| 478 | + tests := []struct { |
| 479 | + name string |
| 480 | + isAccount bool |
| 481 | + primaryErr *apierr.APIError |
| 482 | + meStatus int |
| 483 | + listStatus int |
| 484 | + accountID string |
| 485 | + wantStatus string |
| 486 | + wantUsername string |
| 487 | + wantAccountID string |
| 488 | + }{ |
| 489 | + { |
| 490 | + name: "workspace check fails, account check succeeds", |
| 491 | + primaryErr: &apierr.APIError{StatusCode: http.StatusBadRequest, Message: "Unable to load OAuth Config"}, |
| 492 | + meStatus: http.StatusNotFound, |
| 493 | + listStatus: http.StatusOK, |
| 494 | + accountID: "test-acct", |
| 495 | + wantStatus: "success", |
| 496 | + wantAccountID: "test-acct", |
| 497 | + }, |
| 498 | + { |
| 499 | + name: "no second call without an account id", |
| 500 | + primaryErr: &apierr.APIError{StatusCode: http.StatusBadRequest, Message: "Unable to load OAuth Config"}, |
| 501 | + wantStatus: "error", |
| 502 | + }, |
| 503 | + { |
| 504 | + name: "account check fails, workspace check succeeds", |
| 505 | + isAccount: true, |
| 506 | + primaryErr: &apierr.APIError{StatusCode: http.StatusForbidden, Message: "This API is disabled for users without account admin status"}, |
| 507 | + meStatus: http.StatusOK, |
| 508 | + listStatus: http.StatusNotFound, |
| 509 | + accountID: "test-acct", |
| 510 | + wantStatus: "success", |
| 511 | + wantUsername: "fallback-user", |
| 512 | + }, |
| 513 | + { |
| 514 | + name: "workspace branch, both checks fail", |
| 515 | + primaryErr: &apierr.APIError{StatusCode: http.StatusBadRequest, Message: "Unable to load OAuth Config"}, |
| 516 | + meStatus: http.StatusNotFound, |
| 517 | + listStatus: http.StatusUnauthorized, |
| 518 | + accountID: "test-acct", |
| 519 | + wantStatus: "error", |
| 520 | + }, |
| 521 | + { |
| 522 | + name: "account branch, both checks fail", |
| 523 | + isAccount: true, |
| 524 | + primaryErr: &apierr.APIError{StatusCode: http.StatusUnauthorized, Message: "credentials expired"}, |
| 525 | + meStatus: http.StatusUnauthorized, |
| 526 | + listStatus: http.StatusNotFound, |
| 527 | + accountID: "test-acct", |
| 528 | + wantStatus: "error", |
| 529 | + }, |
| 530 | + } |
| 531 | + |
| 532 | + for _, tc := range tests { |
| 533 | + t.Run(tc.name, func(t *testing.T) { |
| 534 | + t.Setenv("DATABRICKS_ACCOUNT_ID", "") |
| 535 | + server := describeVerifyServer(t, tc.accountID, tc.meStatus, tc.listStatus) |
| 536 | + cfg := &config.Config{} |
| 537 | + var cmd *cobra.Command |
| 538 | + if tc.isAccount { |
| 539 | + cmd = newDescribeAccountCmd(t, cfg, tc.primaryErr) |
| 540 | + } else { |
| 541 | + cmd = newDescribeWorkspaceCmd(t, cfg, tc.primaryErr) |
| 542 | + } |
| 543 | + attrs := map[string]string{ |
| 544 | + "host": server.URL, |
| 545 | + "token": "test-token", |
| 546 | + "auth_type": "pat", |
| 547 | + } |
| 548 | + if tc.accountID != "" { |
| 549 | + attrs["account_id"] = tc.accountID |
| 550 | + } |
| 551 | + |
| 552 | + status, err := getAuthStatus(cmd, []string{}, false, resolveCfg(t, cfg, attrs, tc.isAccount)) |
| 553 | + require.NoError(t, err) |
| 554 | + require.Equal(t, tc.wantStatus, status.Status) |
| 555 | + assert.Equal(t, tc.wantUsername, status.Username) |
| 556 | + assert.Equal(t, tc.wantAccountID, status.AccountID) |
| 557 | + if tc.wantStatus == "error" { |
| 558 | + assert.ErrorIs(t, status.Error, tc.primaryErr) |
| 559 | + } |
| 560 | + }) |
| 561 | + } |
| 562 | +} |
| 563 | + |
400 | 564 | func TestGetWorkspaceAuthStatus_NonU2M_OmitsTokenStorage(t *testing.T) { |
401 | 565 | ctx := t.Context() |
402 | 566 | m := mocks.NewMockWorkspaceClient(t) |
|
0 commit comments