Skip to content

Commit 22a957e

Browse files
authored
auth: add unit and acceptance tests for secure-storage guards (#5077)
## Why I smoke-tested the experimental secure token storage opt-in against a live Databricks workspace (env-var and config-file paths) and noticed three gaps in automated coverage that CI does not catch today: 1. The cmd-layer guard that keeps secure, plaintext, and unknown modes from writing legacy host-keyed entries has no direct unit test. 2. The config-file source of the storage mode has no invalid-value acceptance test (only the env var does). 3. Nothing asserts that \`DATABRICKS_AUTH_STORAGE\` beats \`[__settings__] auth_storage\` in \`.databrickscfg\`. All three are safe to automate without touching a real OS keyring. ## Changes **Before:** \`dualWriteLegacyHostKey\` is only exercised indirectly through login-flow tests; storage-mode acceptance coverage is limited to \`invalid-env\` and \`legacy-env-default\`. **Now:** three test additions, no behavior change. - \`cmd/auth/login_test.go\`: add \`TestDualWriteLegacyHostKey\` with five table-driven cases (legacy mirrors host key, legacy no-op on empty cache, secure / plaintext / unknown modes all skip dual-write). - \`acceptance/cmd/auth/storage-modes/invalid-config/\`: mirror of \`invalid-env\` for the config-file source. Writes \`[__settings__] auth_storage = bogus\` and asserts the error names \`auth_storage\` as the source. - \`acceptance/cmd/auth/storage-modes/env-overrides-config/\`: config says \`secure\`, env says \`legacy\`; \`auth logout\` clears the file-backed cache, proving the env override wins and the keyring is never touched. ## Test plan - [x] \`go test ./cmd/auth/... ./libs/auth/...\` passes - [x] \`go test ./acceptance -run 'TestAccept/cmd/auth/'\` passes - [x] \`make checks\` clean - [x] \`make lint\` 0 issues - [x] New unit test passes in isolation and with \`-race\` - [x] Both acceptance tests pass under both \`DATABRICKS_BUNDLE_ENGINE\` matrix variants Safe for CI: none of the added tests touch a real OS keyring, per the secure-storage rollout plan.
1 parent 08c97f8 commit 22a957e

7 files changed

Lines changed: 140 additions & 0 deletions

File tree

acceptance/cmd/auth/storage-modes/env-overrides-config/out.test.toml

Lines changed: 5 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
2+
=== Token cache keys before logout
3+
[
4+
"dev"
5+
]
6+
7+
>>> [CLI] auth logout --profile dev --auto-approve
8+
Logged out of profile "dev". Use --delete to also remove it from the config file.
9+
10+
=== Token cache keys after logout (should be empty)
11+
[]
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
export DATABRICKS_AUTH_STORAGE=legacy
2+
3+
cat > "./home/.databrickscfg" <<ENDCFG
4+
[__settings__]
5+
auth_storage = secure
6+
7+
[dev]
8+
host = https://accounts.cloud.databricks.com
9+
account_id = account-id
10+
auth_type = databricks-cli
11+
ENDCFG
12+
13+
mkdir -p "./home/.databricks"
14+
cat > "./home/.databricks/token-cache.json" <<ENDCACHE
15+
{
16+
"version": 1,
17+
"tokens": {
18+
"dev": {
19+
"access_token": "dev-cached-token",
20+
"token_type": "Bearer"
21+
}
22+
}
23+
}
24+
ENDCACHE
25+
26+
title "Token cache keys before logout\n"
27+
jq -S '.tokens | keys' "./home/.databricks/token-cache.json"
28+
29+
# Env var = legacy must beat [__settings__] auth_storage = secure.
30+
# Proof: logout clears the file-backed token-cache.json. If the env override
31+
# were ignored, the CLI would try the keyring and leave the file cache alone.
32+
trace $CLI auth logout --profile dev --auto-approve
33+
34+
title "Token cache keys after logout (should be empty)\n"
35+
jq -S '.tokens | keys' "./home/.databricks/token-cache.json"

acceptance/cmd/auth/storage-modes/invalid-config/out.test.toml

Lines changed: 5 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
2+
>>> [CLI] auth token --profile nonexistent
3+
Error: auth_storage: unknown storage mode "bogus" (want legacy, secure, or plaintext)
4+
5+
Exit code: 1
Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
cat > "./home/.databrickscfg" <<ENDCFG
2+
[__settings__]
3+
auth_storage = bogus
4+
ENDCFG
5+
6+
# auth token is the smallest reproducer that resolves the storage mode
7+
# without performing any network I/O.
8+
trace $CLI auth token --profile nonexistent

cmd/auth/login_test.go

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"time"
1616

1717
"github.com/databricks/cli/libs/auth"
18+
"github.com/databricks/cli/libs/auth/storage"
1819
"github.com/databricks/cli/libs/cmdio"
1920
"github.com/databricks/cli/libs/databrickscfg/profile"
2021
"github.com/databricks/cli/libs/env"
@@ -1089,3 +1090,73 @@ func TestLoginRejectsPositionalArgWithProfileFlag(t *testing.T) {
10891090
err := cmd.Execute()
10901091
assert.ErrorContains(t, err, `argument "https://example.com" cannot be combined with --host or --profile`)
10911092
}
1093+
1094+
func TestDualWriteLegacyHostKey(t *testing.T) {
1095+
const (
1096+
profileName = "dual-profile"
1097+
host = "https://dual-host.example.com"
1098+
)
1099+
tok := &oauth2.Token{AccessToken: "abc", RefreshToken: "r"}
1100+
1101+
cacheWithToken := func() *inMemoryTokenCache {
1102+
return &inMemoryTokenCache{Tokens: map[string]*oauth2.Token{profileName: tok}}
1103+
}
1104+
emptyCache := func() *inMemoryTokenCache {
1105+
return &inMemoryTokenCache{Tokens: map[string]*oauth2.Token{}}
1106+
}
1107+
newArg := func(t *testing.T) *u2m.BasicDiscoveryOAuthArgument {
1108+
arg, err := u2m.NewBasicDiscoveryOAuthArgument(profileName)
1109+
require.NoError(t, err)
1110+
arg.SetDiscoveredHost(host)
1111+
return arg
1112+
}
1113+
1114+
cases := []struct {
1115+
name string
1116+
mode storage.StorageMode
1117+
cache func() *inMemoryTokenCache
1118+
wantHostKey bool
1119+
}{
1120+
{
1121+
name: "legacy mirrors cached token under host key",
1122+
mode: storage.StorageModeLegacy,
1123+
cache: cacheWithToken,
1124+
wantHostKey: true,
1125+
},
1126+
{
1127+
name: "legacy is a no-op when cache has no entry",
1128+
mode: storage.StorageModeLegacy,
1129+
cache: emptyCache,
1130+
},
1131+
{
1132+
name: "secure skips dual-write",
1133+
mode: storage.StorageModeSecure,
1134+
cache: cacheWithToken,
1135+
},
1136+
{
1137+
name: "plaintext skips dual-write",
1138+
mode: storage.StorageModePlaintext,
1139+
cache: cacheWithToken,
1140+
},
1141+
{
1142+
name: "unknown mode skips dual-write",
1143+
mode: storage.StorageModeUnknown,
1144+
cache: cacheWithToken,
1145+
},
1146+
}
1147+
1148+
for _, tc := range cases {
1149+
t.Run(tc.name, func(t *testing.T) {
1150+
c := tc.cache()
1151+
dualWriteLegacyHostKey(t.Context(), c, newArg(t), tc.mode)
1152+
1153+
got, err := c.Lookup(host)
1154+
if tc.wantHostKey {
1155+
require.NoError(t, err)
1156+
assert.Equal(t, tok, got)
1157+
} else {
1158+
assert.ErrorIs(t, err, cache.ErrNotFound)
1159+
}
1160+
})
1161+
}
1162+
}

0 commit comments

Comments
 (0)