Skip to content

Commit 5f3cb69

Browse files
samgst-amazonsamgst-amazon
authored
Add source to load registration metrics (#5075)
* initial commit * add source to loading metrics * detekt * fix tests * fix more tests
1 parent 1b6357a commit 5f3cb69

File tree

6 files changed

+41
-34
lines changed

6 files changed

+41
-34
lines changed

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCache.kt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -101,15 +101,15 @@ class DiskCache(
101101
clientRegistrationCache(ssoRegion).tryDeleteIfExists()
102102
}
103103

104-
override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey): ClientRegistration? {
105-
LOG.info { "loadClientRegistration for $cacheKey" }
104+
override fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey, source: String): ClientRegistration? {
105+
LOG.info { "loadClientRegistration:$source for $cacheKey" }
106106
val inputStream = clientRegistrationCache(cacheKey).tryInputStreamIfExists()
107107
if (inputStream == null) {
108108
val stage = LoadCredentialStage.ACCESS_FILE
109109
LOG.info { "Failed to load Client Registration: cache file does not exist" }
110110
AuthTelemetry.modifyConnection(
111111
action = "Load cache file",
112-
source = "loadClientRegistration",
112+
source = "loadClientRegistration:$source",
113113
result = Result.Failed,
114114
reason = "Failed to load Client Registration",
115115
reasonDesc = "Load Step:$stage failed. Cache file does not exist"

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ class SsoAccessTokenProvider(
194194

195195
@Deprecated("Device authorization grant flow is deprecated")
196196
private fun registerDAGClient(): ClientRegistration {
197-
loadDagClientRegistration()?.let {
197+
loadDagClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT.toString())?.let {
198198
return it
199199
}
200200

@@ -235,7 +235,7 @@ class SsoAccessTokenProvider(
235235
}
236236

237237
private fun registerPkceClient(): PKCEClientRegistration {
238-
loadPkceClientRegistration()?.let {
238+
loadPkceClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT.toString())?.let {
239239
return it
240240
}
241241

@@ -431,8 +431,8 @@ class SsoAccessTokenProvider(
431431
stageName = RefreshCredentialStage.LOAD_REGISTRATION
432432
val registration = try {
433433
when (currentToken) {
434-
is DeviceAuthorizationGrantToken -> loadDagClientRegistration()
435-
is PKCEAuthorizationGrantToken -> loadPkceClientRegistration()
434+
is DeviceAuthorizationGrantToken -> loadDagClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString())
435+
is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString())
436436
}
437437
} catch (e: Exception) {
438438
val message = e.message ?: "$stageName: ${e::class.java.name}"
@@ -505,6 +505,11 @@ class SsoAccessTokenProvider(
505505
}
506506
}
507507

508+
enum class SourceOfLoadRegistration {
509+
REGISTER_CLIENT,
510+
REFRESH_TOKEN,
511+
}
512+
508513
private enum class RefreshCredentialStage {
509514
VALIDATE_REFRESH_TOKEN,
510515
LOAD_REGISTRATION,
@@ -514,13 +519,13 @@ class SsoAccessTokenProvider(
514519
SAVE_TOKEN,
515520
}
516521

517-
private fun loadDagClientRegistration(): ClientRegistration? =
518-
cache.loadClientRegistration(dagClientRegistrationCacheKey)?.let {
522+
private fun loadDagClientRegistration(source: String): ClientRegistration? =
523+
cache.loadClientRegistration(dagClientRegistrationCacheKey, source)?.let {
519524
return it
520525
}
521526

522-
private fun loadPkceClientRegistration(): PKCEClientRegistration? =
523-
cache.loadClientRegistration(pkceClientRegistrationCacheKey)?.let {
527+
private fun loadPkceClientRegistration(source: String): PKCEClientRegistration? =
528+
cache.loadClientRegistration(pkceClientRegistrationCacheKey, source)?.let {
524529
return it as PKCEClientRegistration
525530
}
526531

‎plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoCache.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ interface SsoCache {
77
fun invalidateClientRegistration(ssoRegion: String)
88
fun invalidateAccessToken(ssoUrl: String)
99

10-
fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey): ClientRegistration?
10+
fun loadClientRegistration(cacheKey: ClientRegistrationCacheKey, source: String): ClientRegistration?
1111
fun saveClientRegistration(cacheKey: ClientRegistrationCacheKey, registration: ClientRegistration)
1212
fun invalidateClientRegistration(cacheKey: ClientRegistrationCacheKey)
1313

‎plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/DiskCacheTest.kt‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,8 @@ class DiskCacheTest {
5757
startUrl = ssoUrl,
5858
scopes = scopes,
5959
region = ssoRegion
60-
)
60+
),
61+
"testSource"
6162
)
6263
).isNull()
6364
}
@@ -71,7 +72,7 @@ class DiskCacheTest {
7172
)
7273
cacheLocation.resolve("223224b6f0b4702c1a984be8284fe2c9d9718759.json").writeText("badData")
7374

74-
assertThat(sut.loadClientRegistration(key)).isNull()
75+
assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
7576
}
7677

7778
@Test
@@ -91,7 +92,7 @@ class DiskCacheTest {
9192
""".trimIndent()
9293
)
9394

94-
assertThat(sut.loadClientRegistration(key)).isNull()
95+
assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
9596
}
9697

9798
@Test
@@ -112,7 +113,7 @@ class DiskCacheTest {
112113
""".trimIndent()
113114
)
114115

115-
assertThat(sut.loadClientRegistration(key)).isNull()
116+
assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
116117
}
117118

118119
@Test
@@ -134,7 +135,7 @@ class DiskCacheTest {
134135
""".trimIndent()
135136
)
136137

137-
assertThat(sut.loadClientRegistration(key))
138+
assertThat(sut.loadClientRegistration(key, "testSource"))
138139
.usingRecursiveComparison()
139140
.isEqualTo(
140141
DeviceAuthorizationClientRegistration(
@@ -217,7 +218,7 @@ class DiskCacheTest {
217218
""".trimIndent()
218219
)
219220

220-
assertThat(sut.loadClientRegistration(key))
221+
assertThat(sut.loadClientRegistration(key, "testSource"))
221222
.usingRecursiveComparison()
222223
.isEqualTo(
223224
PKCEClientRegistration(
@@ -323,10 +324,10 @@ class DiskCacheTest {
323324
)
324325
)
325326

326-
assertThat(sut.loadClientRegistration(key1))
327+
assertThat(sut.loadClientRegistration(key1, "testSource"))
327328
.usingRecursiveComparison()
328329
.isEqualTo(
329-
sut.loadClientRegistration(key2)
330+
sut.loadClientRegistration(key2, "testSource")
330331
)
331332
}
332333

@@ -350,11 +351,11 @@ class DiskCacheTest {
350351
region = ssoRegion
351352
)
352353

353-
assertThat(sut.loadClientRegistration(key)).isNotNull()
354+
assertThat(sut.loadClientRegistration(key, "testSource")).isNotNull()
354355

355356
sut.invalidateClientRegistration(key)
356357

357-
assertThat(sut.loadClientRegistration(key)).isNull()
358+
assertThat(sut.loadClientRegistration(key, "testSource")).isNull()
358359
assertThat(cacheFile).doesNotExist()
359360
}
360361

@@ -619,7 +620,7 @@ class DiskCacheTest {
619620
registration.setPosixFilePermissions(emptySet())
620621
assertPosixPermissions(registration, "---------")
621622

622-
assertThat(sut.loadClientRegistration(key)).isNotNull()
623+
assertThat(sut.loadClientRegistration(key, "testSource")).isNotNull()
623624

624625
assertPosixPermissions(registration, "rw-------")
625626
}

‎plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProviderTest.kt‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ class SsoAccessTokenProviderTest {
124124
verify(ssoOidcClient).startDeviceAuthorization(any<StartDeviceAuthorizationRequest>())
125125
verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
126126
verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
127-
verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion })
127+
verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any<String>())
128128
verify(ssoCache).saveAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl }, eq(accessToken))
129129
}
130130

@@ -170,7 +170,7 @@ class SsoAccessTokenProviderTest {
170170
verify(ssoOidcClient).startDeviceAuthorization(any<StartDeviceAuthorizationRequest>())
171171
verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
172172
verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
173-
verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion })
173+
verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion }, any<String>())
174174
verify(ssoCache).saveClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion }, any())
175175
verify(ssoCache).saveAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl }, eq(accessToken))
176176
}
@@ -267,7 +267,7 @@ class SsoAccessTokenProviderTest {
267267
verify(ssoOidcClient).startDeviceAuthorization(any<StartDeviceAuthorizationRequest>())
268268
verify(ssoOidcClient, times(2)).createToken(any<CreateTokenRequest>())
269269
verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
270-
verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion })
270+
verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion }, any<String>())
271271
verify(ssoCache).saveAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl }, eq(accessToken))
272272
}
273273

@@ -296,7 +296,7 @@ class SsoAccessTokenProviderTest {
296296
val refreshedToken = runBlocking { sut.refreshToken(sut.accessToken()) }
297297

298298
verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
299-
verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion })
299+
verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any<String>())
300300
verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
301301
verify(ssoCache).saveAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl }, eq(refreshedToken))
302302
}
@@ -342,7 +342,7 @@ class SsoAccessTokenProviderTest {
342342
)
343343

344344
on(
345-
ssoCache.loadClientRegistration(any<PKCEClientRegistrationCacheKey>())
345+
ssoCache.loadClientRegistration(any<PKCEClientRegistrationCacheKey>(), any<String>())
346346
).thenReturn(
347347
PKCEClientRegistration(
348348
clientType = "public",
@@ -369,7 +369,7 @@ class SsoAccessTokenProviderTest {
369369
val refreshedToken = runBlocking { sut.refreshToken(sut.accessToken()) }
370370

371371
verify(ssoCache).loadAccessToken(any<PKCEAccessTokenCacheKey>())
372-
verify(ssoCache).loadClientRegistration(any<PKCEClientRegistrationCacheKey>())
372+
verify(ssoCache).loadClientRegistration(any<PKCEClientRegistrationCacheKey>(), any<String>())
373373
verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
374374
verify(ssoCache).saveAccessToken(any<PKCEAccessTokenCacheKey>(), eq(refreshedToken))
375375
}
@@ -390,7 +390,7 @@ class SsoAccessTokenProviderTest {
390390
verify(ssoOidcClient).startDeviceAuthorization(any<StartDeviceAuthorizationRequest>())
391391
verify(ssoOidcClient).createToken(any<CreateTokenRequest>())
392392
verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
393-
verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion })
393+
verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion }, any<String>())
394394
}
395395

396396
@Test
@@ -432,7 +432,7 @@ class SsoAccessTokenProviderTest {
432432
verify(ssoOidcClient).startDeviceAuthorization(any<StartDeviceAuthorizationRequest>())
433433
verify(ssoOidcClient, times(2)).createToken(any<CreateTokenRequest>())
434434
verify(ssoCache).loadAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl })
435-
verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion })
435+
verify(ssoCache).loadClientRegistration(argThat<DeviceAuthorizationClientRegistrationCacheKey> { region == ssoRegion }, any<String>())
436436
verify(ssoCache).saveAccessToken(argThat<DeviceGrantAccessTokenCacheKey> { startUrl == ssoUrl }, eq(accessToken))
437437
}
438438

@@ -452,7 +452,7 @@ class SsoAccessTokenProviderTest {
452452

453453
verify(ssoOidcClient).registerClient(any<RegisterClientRequest>())
454454
verify(ssoCache).loadAccessToken(any())
455-
verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion })
455+
verify(ssoCache).loadClientRegistration(argThat { region == ssoRegion }, any<String>())
456456
}
457457

458458
@Test
@@ -492,7 +492,7 @@ class SsoAccessTokenProviderTest {
492492
)
493493

494494
on(
495-
ssoCache.loadClientRegistration(argThat { region == ssoRegion })
495+
ssoCache.loadClientRegistration(argThat { region == ssoRegion }, any<String>())
496496
).thenReturn(
497497
returnValue
498498
)

‎plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/credentials/sso/bearer/InteractiveBearerTokenProviderTest.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import org.junit.jupiter.api.assertThrows
1616
import org.mockito.Mockito
1717
import org.mockito.kotlin.any
1818
import org.mockito.kotlin.argThat
19+
import org.mockito.kotlin.eq
1920
import org.mockito.kotlin.mock
2021
import org.mockito.kotlin.spy
2122
import org.mockito.kotlin.times
@@ -273,7 +274,7 @@ class InteractiveBearerTokenProviderTest {
273274
)
274275

275276
private fun stubClientRegistration() {
276-
whenever(diskCache.loadClientRegistration(any<DeviceAuthorizationClientRegistrationCacheKey>())).thenReturn(
277+
whenever(diskCache.loadClientRegistration(any<DeviceAuthorizationClientRegistrationCacheKey>(), eq("testSource"))).thenReturn(
277278
DeviceAuthorizationClientRegistration(
278279
"",
279280
"",

0 commit comments

Comments
 (0)