2023-11-23 09:04:54 +00:00
|
|
|
package auth
|
|
|
|
|
|
|
|
import (
|
|
|
|
"context"
|
|
|
|
"testing"
|
|
|
|
|
|
|
|
"github.com/databricks/cli/libs/auth"
|
2024-08-14 13:01:00 +00:00
|
|
|
"github.com/databricks/cli/libs/cmdio"
|
2023-11-23 09:04:54 +00:00
|
|
|
"github.com/databricks/cli/libs/env"
|
|
|
|
"github.com/stretchr/testify/assert"
|
2024-08-14 13:01:00 +00:00
|
|
|
"github.com/stretchr/testify/require"
|
2023-11-23 09:04:54 +00:00
|
|
|
)
|
|
|
|
|
|
|
|
func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) {
|
|
|
|
ctx := context.Background()
|
|
|
|
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./imaginary-file/databrickscfg")
|
Improve token refresh flow (#1434)
## Changes
Currently, there are a number of issues with the non-happy-path flows
for token refresh in the CLI.
If the token refresh fails, the raw error message is presented to the
user, as seen below. This message is very difficult for users to
interpret and doesn't give any clear direction on how to resolve this
issue.
```
Error: token refresh: Post "https://adb-<WSID>.azuredatabricks.net/oidc/v1/token": http 400: {"error":"invalid_request","error_description":"Refresh token is invalid"}
```
When logging in again, I've noticed that the timeout for logging in is
very short, only 45 seconds. If a user is using a password manager and
needs to login to that first, or needs to do MFA, 45 seconds may not be
enough time. to an account-level profile, it is quite frustrating for
users to need to re-enter account ID information when that information
is already stored in the user's `.databrickscfg` file.
This PR tackles these two issues. First, the presentation of error
messages from `databricks auth token` is improved substantially by
converting the `error` into a human-readable message. When the refresh
token is invalid, it will present a command for the user to run to
reauthenticate. If the token fetching failed for some other reason, that
reason will be presented in a nice way, providing front-line debugging
steps and ultimately redirecting users to file a ticket at this repo if
they can't resolve the issue themselves. After this PR, the new error
message is:
```
Error: a new access token could not be retrieved because the refresh token is invalid. To reauthenticate, run `.databricks/databricks auth login --host https://adb-<WSID>.azuredatabricks.net`
```
To improve the login flow, this PR modifies `databricks auth login` to
auto-complete the account ID from the profile when present.
Additionally, it increases the login timeout from 45 seconds to 1 hour
to give the user sufficient time to login as needed.
To test this change, I needed to refactor some components of the CLI
around profile management, the token cache, and the API client used to
fetch OAuth tokens. These are now settable in the context, and a
demonstration of how they can be set and used is found in
`auth_test.go`.
Separately, this also demonstrates a sort-of integration test of the CLI
by executing the Cobra command for `databricks auth token` from tests,
which may be useful for testing other end-to-end functionality in the
CLI. In particular, I believe this is necessary in order to set flag
values (like the `--profile` flag in this case) for use in testing.
## Tests
Unit tests cover the unhappy and happy paths using the mocked API
client, token cache, and profiler.
Manually tested
---------
Co-authored-by: Pieter Noordhuis <pieter.noordhuis@databricks.com>
2024-05-16 10:22:09 +00:00
|
|
|
err := setHostAndAccountId(ctx, "foo", &auth.PersistentAuth{Host: "test"}, []string{})
|
2023-11-23 09:04:54 +00:00
|
|
|
assert.NoError(t, err)
|
|
|
|
}
|
2024-08-14 13:01:00 +00:00
|
|
|
|
|
|
|
func TestSetHost(t *testing.T) {
|
|
|
|
var persistentAuth auth.PersistentAuth
|
|
|
|
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
|
|
|
|
ctx, _ := cmdio.SetupTest(context.Background())
|
|
|
|
|
|
|
|
// Test error when both flag and argument are provided
|
|
|
|
persistentAuth.Host = "val from --host"
|
|
|
|
err := setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{"val from [HOST]"})
|
|
|
|
assert.EqualError(t, err, "please only provide a host as an argument or a flag, not both")
|
|
|
|
|
|
|
|
// Test setting host from flag
|
|
|
|
persistentAuth.Host = "val from --host"
|
|
|
|
err = setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{})
|
|
|
|
assert.NoError(t, err)
|
|
|
|
assert.Equal(t, "val from --host", persistentAuth.Host)
|
|
|
|
|
|
|
|
// Test setting host from argument
|
|
|
|
persistentAuth.Host = ""
|
|
|
|
err = setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{"val from [HOST]"})
|
|
|
|
assert.NoError(t, err)
|
|
|
|
assert.Equal(t, "val from [HOST]", persistentAuth.Host)
|
|
|
|
|
|
|
|
// Test setting host from profile
|
|
|
|
persistentAuth.Host = ""
|
|
|
|
err = setHostAndAccountId(ctx, "profile-1", &persistentAuth, []string{})
|
|
|
|
assert.NoError(t, err)
|
|
|
|
assert.Equal(t, "https://www.host1.com", persistentAuth.Host)
|
|
|
|
|
|
|
|
// Test setting host from profile
|
|
|
|
persistentAuth.Host = ""
|
|
|
|
err = setHostAndAccountId(ctx, "profile-2", &persistentAuth, []string{})
|
|
|
|
assert.NoError(t, err)
|
|
|
|
assert.Equal(t, "https://www.host2.com", persistentAuth.Host)
|
|
|
|
|
|
|
|
// Test host is not set. Should prompt.
|
|
|
|
persistentAuth.Host = ""
|
|
|
|
err = setHostAndAccountId(ctx, "", &persistentAuth, []string{})
|
|
|
|
assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify a host using --host")
|
|
|
|
}
|
|
|
|
|
|
|
|
func TestSetAccountId(t *testing.T) {
|
|
|
|
var persistentAuth auth.PersistentAuth
|
|
|
|
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
|
|
|
|
ctx, _ := cmdio.SetupTest(context.Background())
|
|
|
|
|
|
|
|
// Test setting account-id from flag
|
|
|
|
persistentAuth.AccountID = "val from --account-id"
|
|
|
|
err := setHostAndAccountId(ctx, "account-profile", &persistentAuth, []string{})
|
|
|
|
assert.NoError(t, err)
|
|
|
|
assert.Equal(t, "https://accounts.cloud.databricks.com", persistentAuth.Host)
|
|
|
|
assert.Equal(t, "val from --account-id", persistentAuth.AccountID)
|
|
|
|
|
|
|
|
// Test setting account_id from profile
|
|
|
|
persistentAuth.AccountID = ""
|
|
|
|
err = setHostAndAccountId(ctx, "account-profile", &persistentAuth, []string{})
|
|
|
|
require.NoError(t, err)
|
|
|
|
assert.Equal(t, "https://accounts.cloud.databricks.com", persistentAuth.Host)
|
|
|
|
assert.Equal(t, "id-from-profile", persistentAuth.AccountID)
|
|
|
|
|
|
|
|
// Neither flag nor profile account-id is set, should prompt
|
|
|
|
persistentAuth.AccountID = ""
|
|
|
|
persistentAuth.Host = "https://accounts.cloud.databricks.com"
|
|
|
|
err = setHostAndAccountId(ctx, "", &persistentAuth, []string{})
|
|
|
|
assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify an account ID using --account-id")
|
|
|
|
}
|