Repository: snowflakedb/gosnowflake
Pull request #1299 (2025-02-04 12:10:44)

SNOW-1825790 Token cache refactor - v2
### Description

SNOW-XXX Please explain the changes you made here.

### Checklist
- [x] Created tests which fail without the change (if possible)
- [x] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1299?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
Attention: Patch coverage is `88.60759%` with `9 lines` in your changes missing coverage. Please review.
> Project coverage is 82.17%. Comparing base [(`e926883`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/e926883b12016eb30ae25c2027f2cf70714eb2c5?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) to head [(`3a3de37`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/3a3de37b29ef6f3359f12af672f5ada96baaa0ec?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).

| [Files with missing lines](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1299?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Patch % | Lines |
|---|---|---|
| [secure\_storage\_manager.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1299?src=pr&el=tree&filepath=secure_storage_manager.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-c2VjdXJlX3N0b3JhZ2VfbWFuYWdlci5nbw==) | 87.67% | [9 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1299?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #1299      +/-   ##
==========================================
- Coverage   82.20%   82.17%   -0.04%     
==========================================
  Files          55       55              
  Lines       13636    13646      +10     
==========================================
+ Hits        11209    11213       +4     
- Misses       2427     2433       +6     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1299?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).


Patch:
diff --git a/auth.go b/auth.go
index a95b968..ae214bb 100644
--- a/auth.go
+++ b/auth.go
@@ -26,8 +26,6 @@ const (
 )
 
 const (
-	idToken                        = "ID_TOKEN"
-	mfaToken                       = "MFATOKEN"
 	clientStoreTemporaryCredential = "CLIENT_STORE_TEMPORARY_CREDENTIAL"
 	clientRequestMfaToken          = "CLIENT_REQUEST_MFA_TOKEN"
 	idTokenAuthenticator           = "ID_TOKEN"
@@ -365,10 +363,10 @@ func authenticate(
 		logger.WithContext(ctx).Errorln("Authentication FAILED")
 		sc.rest.TokenAccessor.SetTokens("", "", -1)
 		if sessionParameters[clientRequestMfaToken] == true {
-			credentialsStorage.deleteCredential(sc, mfaToken)
+			credentialsStorage.deleteCredential(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))
 		}
 		if sessionParameters[clientStoreTemporaryCredential] == true {
-			credentialsStorage.deleteCredential(sc, idToken)
+			credentialsStorage.deleteCredential(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
 		}
 		code, err := strconv.Atoi(respd.Code)
 		if err != nil {
@@ -384,11 +382,11 @@ func authenticate(
 	sc.rest.TokenAccessor.SetTokens(respd.Data.Token, respd.Data.MasterToken, respd.Data.SessionID)
 	if sessionParameters[clientRequestMfaToken] == true {
 		token := respd.Data.MfaToken
-		credentialsStorage.setCredential(sc, mfaToken, token)
+		credentialsStorage.setCredential(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User), token)
 	}
 	if sessionParameters[clientStoreTemporaryCredential] == true {
 		token := respd.Data.IDToken
-		credentialsStorage.setCredential(sc, idToken, token)
+		credentialsStorage.setCredential(newIDTokenSpec(sc.cfg.Host, sc.cfg.User), token)
 	}
 	return &respd.Data, nil
 }
@@ -523,7 +521,7 @@ func authenticateWithConfig(sc *snowflakeConn) error {
 			sc.cfg.ClientStoreTemporaryCredential = ConfigBoolTrue
 		}
 		if sc.cfg.ClientStoreTemporaryCredential == ConfigBoolTrue {
-			fillCachedIDToken(sc)
+			sc.cfg.IDToken = credentialsStorage.getCredential(newIDTokenSpec(sc.cfg.Host, sc.cfg.User))
 		}
 		// Disable console login by default
 		if sc.cfg.DisableConsoleLogin == configBoolNotSet {
@@ -536,7 +534,7 @@ func authenticateWithConfig(sc *snowflakeConn) error {
 			sc.cfg.ClientRequestMfaToken = ConfigBoolTrue
 		}
 		if sc.cfg.ClientRequestMfaToken == ConfigBoolTrue {
-			fillCachedMfaToken(sc)
+			sc.cfg.MfaToken = credentialsStorage.getCredential(newMfaTokenSpec(sc.cfg.Host, sc.cfg.User))
 		}
 	}
 
@@ -573,11 +571,3 @@ func authenticateWithConfig(sc *snowflakeConn) error {
 	sc.ctx = context.WithValue(sc.ctx, SFSessionIDKey, authData.SessionID)
 	return nil
 }
-
-func fillCachedIDToken(sc *snowflakeConn) {
-	credentialsStorage.getCredential(sc, idToken)
-}
-
-func fillCachedMfaToken(sc *snowflakeConn) {
-	credentialsStorage.getCredential(sc, mfaToken)
-}
diff --git a/secure_storage_manager.go b/secure_storage_manager.go
index fc39266..6be71a5 100644
--- a/secure_storage_manager.go
+++ b/secure_storage_manager.go
@@ -15,16 +15,48 @@ import (
 	"github.com/99designs/keyring"
 )
 
+type tokenType string
+
+const (
+	idToken  tokenType = "ID_TOKEN"
+	mfaToken tokenType = "MFATOKEN"
+)
+
 const (
 	driverName        = "SNOWFLAKE-GO-DRIVER"
 	credCacheDirEnv   = "SF_TEMPORARY_CREDENTIAL_CACHE_DIR"
 	credCacheFileName = "temporary_credential.json"
 )
 
+type secureTokenSpec struct {
+	host, user string
+	tokenType  tokenType
+}
+
+func (t *secureTokenSpec) buildKey() string {
+	return buildCredentialsKey(t.host, t.user, t.tokenType)
+}
+
+func newMfaTokenSpec(host, user string) *secureTokenSpec {
+	return &secureTokenSpec{
+		host,
+		user,
+		mfaToken,
+	}
+}
+
+func newIDTokenSpec(host, user string) *secureTokenSpec {
+	return &secureTokenSpec{
+		host,
+		user,
+		idToken,
+	}
+}
+
 type secureStorageManager interface {
-	setCredential(sc *snowflakeConn, credType, token string)
-	getCredential(sc *snowflakeConn, credType string)
-	deleteCredential(sc *snowflakeConn, credType string)
+	setCredential(tokenSpec *secureTokenSpec, value string)
+	getCredential(tokenSpec *secureTokenSpec) string
+	deleteCredential(tokenSpec *secureTokenSpec)
 }
 
 var credentialsStorage = newSecureStorageManager()
@@ -32,7 +64,12 @@ var credentialsStorage = newSecureStorageManager()
 func newSecureStorageManager() secureStorageManager {
 	switch runtime.GOOS {
 	case "linux":
-		return newFileBasedSecureStorageManager()
+		ssm, err := newFileBasedSecureStorageManager()
+		if err != nil {
+			logger.Debugf("failed to create credentials cache dir. %v", err)
+			return newNoopSecureStorageManager()
+		}
+		return ssm
 	case "darwin", "windows":
 		return newKeyringBasedSecureStorageManager()
 	default:
@@ -46,20 +83,19 @@ type fileBasedSecureStorageManager struct {
 	credCacheLock     sync.RWMutex
 }
 
-func newFileBasedSecureStorageManager() secureStorageManager {
+func newFileBasedSecureStorageManager() (*fileBasedSecureStorageManager, error) {
 	ssm := &fileBasedSecureStorageManager{
 		localCredCache: map[string]string{},
 		credCacheLock:  sync.RWMutex{},
 	}
 	credCacheDir := ssm.buildCredCacheDirPath()
 	if err := ssm.createCacheDir(credCacheDir); err != nil {
-		logger.Debugf("failed to create credentials cache dir. %v", err)
-		return newNoopSecureStorageManager()
+		return nil, err
 	}
 	credCacheFilePath := filepath.Join(credCacheDir, credCacheFileName)
 	logger.Infof("Credentials cache path: %v", credCacheFilePath)
 	ssm.credCacheFilePath = credCacheFilePath
-	return ssm
+	return ssm, nil
 }
 
 func (ssm *fileBasedSecureStorageManager) createCacheDir(credCacheDir string) error {
@@ -87,14 +123,14 @@ func (ssm *fileBasedSecureStorageManager) buildCredCacheDirPath() string {
 	return credCacheDir
 }
 
-func (ssm *fileBasedSecureStorageManager) setCredential(sc *snowflakeConn, credType, token string) {
-	if token == "" {
+func (ssm *fileBasedSecureStorageManager) setCredential(tokenSpec *secureTokenSpec, value string) {
+	if value == "" {
 		logger.Debug("no token provided")
 	} else {
-		credentialsKey := buildCredentialsKey(sc.cfg.Host, sc.cfg.User, credType)
+		credentialsKey := tokenSpec.buildKey()
 		ssm.credCacheLock.Lock()
 		defer ssm.credCacheLock.Unlock()
-		ssm.localCredCache[credentialsKey] = token
+		ssm.localCredCache[credentialsKey] = value
 
 		j, err := json.Marshal(ssm.localCredCache)
 		if err != nil {
@@ -135,8 +171,8 @@ func (ssm *fileBasedSecureStorageManager) setCredential(sc *snowflakeConn, credT
 	}
 }
 
-func (ssm *fileBasedSecureStorageManager) getCredential(sc *snowflakeConn, credType string) {
-	credentialsKey := buildCredentialsKey(sc.cfg.Host, sc.cfg.User, credType)
+func (ssm *fileBasedSecureStorageManager) getCredential(tokenSpec *secureTokenSpec) string {
+	credentialsKey := tokenSpec.buildKey()
 	ssm.credCacheLock.Lock()
 	defer ssm.credCacheLock.Unlock()
 	localCredCache := ssm.readTemporaryCacheFile()
@@ -146,14 +182,7 @@ func (ssm *fileBasedSecureStorageManager) getCredential(sc *snowflakeConn, credT
 	} else {
 		logger.Debug("Returned credential is empty")
 	}
-
-	if credType == idToken {
-		sc.cfg.IDToken = cred
-	} else if credType == mfaToken {
-		sc.cfg.MfaToken = cred
-	} else {
-		logger.Debugf("Unrecognized type %v for local cached credential", credType)
-	}
+	return cred
 }
 
 func (ssm *fileBasedSecureStorageManager) readTemporaryCacheFile() map[string]string {
@@ -171,10 +200,10 @@ func (ssm *fileBasedSecureStorageManager) readTemporaryCacheFile() map[string]st
 	return ssm.localCredCache
 }
 
-func (ssm *fileBasedSecureStorageManager) deleteCredential(sc *snowflakeConn, credType string) {
+func (ssm *fileBasedSecureStorageManager) deleteCredential(tokenSpec *secureTokenSpec) {
 	ssm.credCacheLock.Lock()
 	defer ssm.credCacheLock.Unlock()
-	credentialsKey := buildCredentialsKey(sc.cfg.Host, sc.cfg.User, credType)
+	credentialsKey := tokenSpec.buildKey()
 	delete(ssm.localCredCache, credentialsKey)
 	j, err := json.Marshal(ssm.localCredCache)
 	if err != nil {
@@ -220,37 +249,35 @@ func (ssm *fileBasedSecureStorageManager) writeTemporaryCacheFile(input []byte)
 type keyringSecureStorageManager struct {
 }
 
-func newKeyringBasedSecureStorageManager() secureStorageManager {
+func newKeyringBasedSecureStorageManager() *keyringSecureStorageManager {
 	return &keyringSecureStorageManager{}
 }
 
-func (ssm *keyringSecureStorageManager) setCredential(sc *snowflakeConn, credType, token string) {
-	if token == "" {
+func (ssm *keyringSecureStorageManager) setCredential(tokenSpec *secureTokenSpec, value string) {
+	if value == "" {
 		logger.Debug("no token provided")
 	} else {
-		var credentialsKey string
+		credentialsKey := tokenSpec.buildKey()
 		if runtime.GOOS == "windows" {
-			credentialsKey = driverName + ":" + credType
 			ring, _ := keyring.Open(keyring.Config{
-				WinCredPrefix: strings.ToUpper(sc.cfg.Host),
-				ServiceName:   strings.ToUpper(sc.cfg.User),
+				WinCredPrefix: strings.ToUpper(tokenSpec.host),
+				ServiceName:   strings.ToUpper(tokenSpec.user),
 			})
 			item := keyring.Item{
 				Key:  credentialsKey,
-				Data: []byte(token),
+				Data: []byte(value),
 			}
 			if err := ring.Set(item); err != nil {
 				logger.Debugf("Failed to write to Windows credential manager. Err: %v", err)
 			}
 		} else if runtime.GOOS == "darwin" {
-			credentialsKey = buildCredentialsKey(sc.cfg.Host, sc.cfg.User, credType)
 			ring, _ := keyring.Open(keyring.Config{
 				ServiceName: credentialsKey,
 			})
-			account := strings.ToUpper(sc.cfg.User)
+			account := strings.ToUpper(tokenSpec.user)
 			item := keyring.Item{
 				Key:  account,
-				Data: []byte(token),
+				Data: []byte(value),
 			}
 			if err := ring.Set(item); err != nil {
 				logger.Debugf("Failed to write to keychain. Err: %v", err)
@@ -259,14 +286,13 @@ func (ssm *keyringSecureStorageManager) setCredential(sc *snowflakeConn, credTyp
 	}
 }
 
-func (ssm *keyringSecureStorageManager) getCredential(sc *snowflakeConn, credType string) {
-	var credentialsKey string
+func (ssm *keyringSecureStorageManager) getCredential(tokenSpec *secureTokenSpec) string {
 	cred := ""
+	credentialsKey := tokenSpec.buildKey()
 	if runtime.GOOS == "windows" {
-		credentialsKey = driverName + ":" + credType
 		ring, _ := keyring.Open(keyring.Config{
-			WinCredPrefix: strings.ToUpper(sc.cfg.Host),
-			ServiceName:   strings.ToUpper(sc.cfg.User),
+			WinCredPrefix: strings.ToUpper(tokenSpec.host),
+			ServiceName:   strings.ToUpper(tokenSpec.user),
 		})
 		i, err := ring.Get(credentialsKey)
 		if err != nil {
@@ -274,11 +300,10 @@ func (ssm *keyringSecureStorageManager) getCredential(sc *snowflakeConn, credTyp
 		}
 		cred = string(i.Data)
 	} else if runtime.GOOS == "darwin" {
-		credentialsKey = buildCredentialsKey(sc.cfg.Host, sc.cfg.User, credType)
 		ring, _ := keyring.Open(keyring.Config{
 			ServiceName: credentialsKey,
 		})
-		account := strings.ToUpper(sc.cfg.User)
+		account := strings.ToUpper(tokenSpec.user)
 		i, err := ring.Get(account)
 		if err != nil {
 			logger.Debugf("Failed to find the item in keychain or item does not exist. Error: %v", err)
@@ -290,33 +315,25 @@ func (ssm *keyringSecureStorageManager) getCredential(sc *snowflakeConn, credTyp
 			logger.Debug("Successfully read token. Returning as string")
 		}
 	}
-
-	if credType == idToken {
-		sc.cfg.IDToken = cred
-	} else if credType == mfaToken {
-		sc.cfg.MfaToken = cred
-	} else {
-		logger.Debugf("Unrecognized type %v for local cached credential", credType)
-	}
+	return cred
 }
 
-func (ssm *keyringSecureStorageManager) deleteCredential(sc *snowflakeConn, credType string) {
-	credentialsKey := driverName + ":" + credType
+func (ssm *keyringSecureStorageManager) deleteCredential(tokenSpec *secureTokenSpec) {
+	credentialsKey := tokenSpec.buildKey()
 	if runtime.GOOS == "windows" {
 		ring, _ := keyring.Open(keyring.Config{
-			WinCredPrefix: strings.ToUpper(sc.cfg.Host),
-			ServiceName:   strings.ToUpper(sc.cfg.User),
+			WinCredPrefix: strings.ToUpper(tokenSpec.host),
+			ServiceName:   strings.ToUpper(tokenSpec.user),
 		})
-		err := ring.Remove(credentialsKey)
+		err := ring.Remove(string(credentialsKey))
 		if err != nil {
 			logger.Debugf("Failed to delete credentialsKey in Windows Credential Manager. Error: %v", err)
 		}
 	} else if runtime.GOOS == "darwin" {
-		credentialsKey = buildCredentialsKey(sc.cfg.Host, sc.cfg.User, credType)
 		ring, _ := keyring.Open(keyring.Config{
 			ServiceName: credentialsKey,
 		})
-		account := strings.ToUpper(sc.cfg.User)
+		account := strings.ToUpper(tokenSpec.user)
 		err := ring.Remove(account)
 		if err != nil {
 			logger.Debugf("Failed to delete credentialsKey in keychain. Error: %v", err)
@@ -324,25 +341,26 @@ func (ssm *keyringSecureStorageManager) deleteCredential(sc *snowflakeConn, cred
 	}
 }
 
-func buildCredentialsKey(host, user, credType string) string {
+func buildCredentialsKey(host, user string, credType tokenType) string {
 	host = strings.ToUpper(host)
 	user = strings.ToUpper(user)
-	credType = strings.ToUpper(credType)
-	return host + ":" + user + ":" + driverName + ":" + credType
+	credTypeStr := strings.ToUpper(string(credType))
+	return host + ":" + user + ":" + driverName + ":" + credTypeStr
 }
 
 type noopSecureStorageManager struct {
 }
 
-func newNoopSecureStorageManager() secureStorageManager {
+func newNoopSecureStorageManager() *noopSecureStorageManager {
 	return &noopSecureStorageManager{}
 }
 
-func (ssm *noopSecureStorageManager) setCredential(sc *snowflakeConn, credType, token string) {
+func (ssm *noopSecureStorageManager) setCredential(_ *secureTokenSpec, _ string) {
 }
 
-func (ssm *noopSecureStorageManager) getCredential(sc *snowflakeConn, credType string) {
+func (ssm *noopSecureStorageManager) getCredential(_ *secureTokenSpec) string {
+	return ""
 }
 
-func (ssm *noopSecureStorageManager) deleteCredential(sc *snowflakeConn, credType string) { //TODO implement me
+func (ssm *noopSecureStorageManager) deleteCredential(_ *secureTokenSpec) { //TODO implement me
 }


Test patch:
diff --git a/secure_storage_manager_test.go b/secure_storage_manager_test.go
index 3b3c232..2229970 100644
--- a/secure_storage_manager_test.go
+++ b/secure_storage_manager_test.go
@@ -6,57 +6,22 @@ import (
 	"testing"
 )
 
-type tcTargets struct {
-	host     string
-	user     string
-	credType string
-	out      string
-}
-
-type tcCredentials struct {
-	credType string
-	token    string
-}
-
 func TestSetAndGetCredentialMfa(t *testing.T) {
-	skipOnMac(t, "keyring asks for password")
-	fakeMfaToken := "fakeMfaToken"
-	expectedMfaToken := "fakeMfaToken"
-	sc := getDefaultSnowflakeConn()
-	sc.cfg.Host = "testhost"
-	credentialsStorage.setCredential(sc, mfaToken, fakeMfaToken)
-	credentialsStorage.getCredential(sc, mfaToken)
+	for _, tokenSpec := range []*secureTokenSpec{
+		newMfaTokenSpec("testhost", "testuser"),
+		newIDTokenSpec("testhost", "testuser"),
+	} {
+		t.Run(string(tokenSpec.tokenType), func(t *testing.T) {
+			skipOnMac(t, "keyring asks for password")
+			fakeMfaToken := "test token"
+			tokenSpec := newMfaTokenSpec("testHost", "testUser")
+			credentialsStorage.setCredential(tokenSpec, fakeMfaToken)
+			assertEqualE(t, credentialsStorage.getCredential(tokenSpec), fakeMfaToken)
 
-	if sc.cfg.MfaToken != expectedMfaToken {
-		t.Fatalf("Expected mfa token %v but got %v", expectedMfaToken, sc.cfg.MfaToken)
-	}
-
-	// delete credential and check it no longer exists
-	credentialsStorage.deleteCredential(sc, mfaToken)
-	credentialsStorage.getCredential(sc, mfaToken)
-	if sc.cfg.MfaToken != "" {
-		t.Fatalf("Expected mfa token to be empty but got %v", sc.cfg.MfaToken)
-	}
-}
-
-func TestSetAndGetCredentialIdToken(t *testing.T) {
-	skipOnMac(t, "keyring asks for password")
-	fakeIDToken := "fakeIDToken"
-	expectedIDToken := "fakeIDToken"
-	sc := getDefaultSnowflakeConn()
-	sc.cfg.Host = "testhost"
-	credentialsStorage.setCredential(sc, idToken, fakeIDToken)
-	credentialsStorage.getCredential(sc, idToken)
-
-	if sc.cfg.IDToken != expectedIDToken {
-		t.Fatalf("Expected id token %v but got %v", expectedIDToken, sc.cfg.IDToken)
-	}
-
-	// delete credential and check it no longer exists
-	credentialsStorage.deleteCredential(sc, idToken)
-	credentialsStorage.getCredential(sc, idToken)
-	if sc.cfg.IDToken != "" {
-		t.Fatalf("Expected id token to be empty but got %v", sc.cfg.IDToken)
+			// delete credential and check it no longer exists
+			credentialsStorage.deleteCredential(tokenSpec)
+			assertEqualE(t, credentialsStorage.getCredential(tokenSpec), "")
+		})
 	}
 }
 
@@ -65,38 +30,34 @@ func TestStoreTemporaryCredental(t *testing.T) {
 		t.Skip("cannot write to github file system")
 	}
 
-	testcases := []tcCredentials{
-		{mfaToken, "598ghFnjfh8BBgmf45mmhgkfRR45mgkt5"},
-		{idToken, "090Arftf54Jk3gh57ggrVvf09lJa3DD"},
+	testcases := []struct {
+		tokenSpec *secureTokenSpec
+		value     string
+	}{
+		{newMfaTokenSpec("testhost", "testuser"), "598ghFnjfh8BBgmf45mmhgkfRR45mgkt5"},
+		{newIDTokenSpec("testhost", "testuser"), "090Arftf54Jk3gh57ggrVvf09lJa3DD"},
 	}
 
-	ssm := newFileBasedSecureStorageManager()
-	_, ok := ssm.(*fileBasedSecureStorageManager)
-	assertTrueF(t, ok)
+	ssm, err := newFileBasedSecureStorageManager()
+	assertNilF(t, err)
 
-	sc := getDefaultSnowflakeConn()
 	for _, test := range testcases {
-		t.Run(test.token, func(t *testing.T) {
-			ssm.setCredential(sc, test.credType, test.token)
-			ssm.getCredential(sc, test.credType)
-			if test.credType == mfaToken {
-				assertEqualE(t, sc.cfg.MfaToken, test.token)
-			} else {
-				assertEqualE(t, sc.cfg.IDToken, test.token)
-			}
-			ssm.deleteCredential(sc, test.credType)
-			ssm.getCredential(sc, test.credType)
-			if test.credType == mfaToken {
-				assertEqualE(t, sc.cfg.MfaToken, "")
-			} else {
-				assertEqualE(t, sc.cfg.IDToken, "")
-			}
+		t.Run(test.value, func(t *testing.T) {
+			ssm.setCredential(test.tokenSpec, test.value)
+			assertEqualE(t, ssm.getCredential(test.tokenSpec), test.value)
+			ssm.deleteCredential(test.tokenSpec)
+			assertEqualE(t, ssm.getCredential(test.tokenSpec), "")
 		})
 	}
 }
 
 func TestBuildCredentialsKey(t *testing.T) {
-	testcases := []tcTargets{
+	testcases := []struct {
+		host     string
+		user     string
+		credType tokenType
+		out      string
+	}{
 		{"testaccount.snowflakecomputing.com", "testuser", "mfaToken", "TESTACCOUNT.SNOWFLAKECOMPUTING.COM:TESTUSER:SNOWFLAKE-GO-DRIVER:MFATOKEN"},
 		{"testaccount.snowflakecomputing.com", "testuser", "IdToken", "TESTACCOUNT.SNOWFLAKECOMPUTING.COM:TESTUSER:SNOWFLAKE-GO-DRIVER:IDTOKEN"},
 	}