Skip to content

Commit ab19721

Browse files
xinsong-cuixinsong-cui
committed
style & rename codegen file
1 parent 3853646 commit ab19721

File tree

3 files changed

+106
-117
lines changed

3 files changed

+106
-117
lines changed

‎codegen/aws-sdk-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/customization/EnvironmentBearerTokenCustomization.kt‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@
44
*/
55
package aws.sdk.kotlin.codegen.customization
66

7+
import aws.sdk.kotlin.codegen.SdkIdTransform
78
import aws.sdk.kotlin.codegen.ServiceClientCompanionObjectWriter
9+
import aws.sdk.kotlin.codegen.withTransform
810
import software.amazon.smithy.kotlin.codegen.KotlinSettings
911
import software.amazon.smithy.kotlin.codegen.core.*
1012
import software.amazon.smithy.kotlin.codegen.integration.AppendingSectionWriter
@@ -45,18 +47,16 @@ class EnvironmentBearerTokenCustomization : KotlinIntegration {
4547

4648
override fun writeAdditionalFiles(ctx: CodegenContext, delegator: KotlinDelegator) {
4749
val serviceShape = ctx.model.expectShape<ServiceShape>(ctx.settings.service)
48-
val serviceName = clientName(ctx.settings.sdkId)
4950
val packageName = ctx.settings.pkg.name
5051

5152
delegator.useFileWriter(
52-
"Finalize${serviceName}EnvironmentBearerTokenConfig.kt",
53+
"FinalizeBearerTokenConfig.kt",
5354
"$packageName.auth",
5455
) { writer ->
5556
renderEnvironmentBearerTokenConfig(
5657
writer,
5758
ctx,
5859
serviceShape,
59-
serviceName,
6060
)
6161
}
6262
}
@@ -65,18 +65,17 @@ class EnvironmentBearerTokenCustomization : KotlinIntegration {
6565
writer: KotlinWriter,
6666
ctx: CodegenContext,
6767
serviceShape: ServiceShape,
68-
serviceName: String,
6968
) {
7069
val serviceSymbol = ctx.symbolProvider.toSymbol(serviceShape)
7170
val signingServiceName = AwsSignatureVersion4.signingServiceName(serviceShape)
7271
// Transform signing name to environment variable name
73-
val envVarName = "AWS_BEARER_TOKEN_" + signingServiceName.replace("""[-\s]""".toRegex(), "_").uppercase()
72+
val envVarSuffix = signingServiceName.withTransform(SdkIdTransform.UpperSnakeCase)
73+
val envVarName = "AWS_BEARER_TOKEN_$envVarSuffix"
7474
val authSchemeId = RuntimeTypes.Auth.Identity.AuthSchemeId
7575

7676
writer.withBlock(
77-
"internal fun finalize#1LEnvironmentBearerTokenConfig(builder: #2T.Builder, provider: #3T = #3T.System) {",
77+
"internal fun finalizeBearerTokenConfig(builder: #1T.Builder, provider: #2T = #2T.System) {",
7878
"}",
79-
serviceName,
8079
serviceSymbol,
8180
RuntimeTypes.Core.Utils.PlatformProvider,
8281
) {
@@ -113,7 +112,7 @@ class EnvironmentBearerTokenCustomization : KotlinIntegration {
113112
val serviceName = clientName(writer.getContextValue(ServiceClientGenerator.Sections.CompanionObject.SdkId))
114113

115114
val environmentBearerTokenConfig = buildSymbol {
116-
name = "finalize${serviceName}EnvironmentBearerTokenConfig"
115+
name = "finalizeBearerTokenConfig"
117116
namespace = "aws.sdk.kotlin.services.${serviceName.lowercase()}.auth"
118117
}
119118

‎codegen/aws-sdk-codegen/src/test/kotlin/aws/sdk/kotlin/codegen/customization/EnvironmentBearerTokenCustomizationTest.kt‎

Lines changed: 66 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -10,81 +10,43 @@ import kotlin.test.assertFalse
1010
import kotlin.test.assertTrue
1111

1212
class EnvironmentBearerTokenCustomizationTest {
13-
private val bedrockModel = """
14-
namespace com.test
15-
use aws.auth#sigv4
16-
use aws.api#service
17-
use smithy.api#httpBearerAuth
18-
19-
@sigv4(name: "bedrock")
20-
@httpBearerAuth
21-
@service(sdkId: "Bedrock")
22-
service Bedrock {
23-
version: "1.0.0"
24-
}
25-
""".trimIndent().toSmithyModel()
26-
27-
private val bedrockRuntimeModel = """
28-
namespace com.test
29-
use aws.auth#sigv4
30-
use aws.api#service
31-
use smithy.api#httpBearerAuth
32-
33-
@sigv4(name: "bedrock")
34-
@httpBearerAuth
35-
@service(sdkId: "Bedrock Runtime")
36-
service BedrockRuntime {
37-
version: "1.0.0"
38-
}
39-
""".trimIndent().toSmithyModel()
40-
41-
private val nonBedrockModel = """
42-
namespace com.test
43-
use aws.auth#sigv4
44-
use aws.api#service
45-
use smithy.api#httpBearerAuth
46-
47-
@sigv4(name: "s3")
48-
@httpBearerAuth
49-
@service(sdkId: "S3")
50-
service S3 {
51-
version: "1.0.0"
52-
}
53-
""".trimIndent().toSmithyModel()
54-
55-
private val noSigV4Model = """
56-
namespace com.test
57-
use aws.api#service
58-
use smithy.api#httpBearerAuth
59-
60-
@service(sdkId: "NoSigV4")
61-
@httpBearerAuth
62-
service NoSigV4 {
63-
version: "1.0.0"
64-
}
65-
""".trimIndent().toSmithyModel()
66-
67-
private val noBearerAuthModel = """
68-
namespace com.test
69-
use aws.auth#sigv4
70-
use aws.api#service
71-
72-
@sigv4(name: "bedrock")
73-
@service(sdkId: "BedrockNoBearerAuth")
74-
service BedrockNoBearerAuth {
75-
version: "1.0.0"
76-
}
77-
""".trimIndent().toSmithyModel()
78-
7913
@Test
8014
fun `test customization enabled for bedrock sigv4 signing name`() {
15+
val bedrockModel = """
16+
namespace com.test
17+
use aws.auth#sigv4
18+
use aws.api#service
19+
use smithy.api#httpBearerAuth
20+
21+
@sigv4(name: "bedrock")
22+
@httpBearerAuth
23+
@service(sdkId: "Bedrock")
24+
service Bedrock {
25+
version: "1.0.0"
26+
}
27+
""".trimIndent().toSmithyModel()
28+
8129
assertTrue {
8230
EnvironmentBearerTokenCustomization()
8331
.enabledForService(bedrockModel, bedrockModel.defaultSettings())
8432
}
8533
}
8634

8735
fun `test customization enabled for bedrock sigv4 signing name with different sdkId`() {
36+
val bedrockRuntimeModel = """
37+
namespace com.test
38+
use aws.auth#sigv4
39+
use aws.api#service
40+
use smithy.api#httpBearerAuth
41+
42+
@sigv4(name: "bedrock")
43+
@httpBearerAuth
44+
@service(sdkId: "Bedrock Runtime")
45+
service BedrockRuntime {
46+
version: "1.0.0"
47+
}
48+
""".trimIndent().toSmithyModel()
49+
8850
assertTrue {
8951
EnvironmentBearerTokenCustomization()
9052
.enabledForService(bedrockRuntimeModel, bedrockRuntimeModel.defaultSettings())
@@ -93,6 +55,20 @@ class EnvironmentBearerTokenCustomizationTest {
9355

9456
@Test
9557
fun `test customization not enabled for non-bedrock sigv4 signing name`() {
58+
val nonBedrockModel = """
59+
namespace com.test
60+
use aws.auth#sigv4
61+
use aws.api#service
62+
use smithy.api#httpBearerAuth
63+
64+
@sigv4(name: "s3")
65+
@httpBearerAuth
66+
@service(sdkId: "S3")
67+
service S3 {
68+
version: "1.0.0"
69+
}
70+
""".trimIndent().toSmithyModel()
71+
9672
assertFalse {
9773
EnvironmentBearerTokenCustomization()
9874
.enabledForService(nonBedrockModel, nonBedrockModel.defaultSettings())
@@ -101,6 +77,18 @@ class EnvironmentBearerTokenCustomizationTest {
10177

10278
@Test
10379
fun `test customization not enabled for model without sigv4 trait`() {
80+
val noSigV4Model = """
81+
namespace com.test
82+
use aws.api#service
83+
use smithy.api#httpBearerAuth
84+
85+
@service(sdkId: "NoSigV4")
86+
@httpBearerAuth
87+
service NoSigV4 {
88+
version: "1.0.0"
89+
}
90+
""".trimIndent().toSmithyModel()
91+
10492
assertFalse {
10593
EnvironmentBearerTokenCustomization()
10694
.enabledForService(noSigV4Model, noSigV4Model.defaultSettings())
@@ -109,6 +97,18 @@ class EnvironmentBearerTokenCustomizationTest {
10997

11098
@Test
11199
fun `test customization not enabled for model without bearer auth trait`() {
100+
val noBearerAuthModel = """
101+
namespace com.test
102+
use aws.auth#sigv4
103+
use aws.api#service
104+
105+
@sigv4(name: "bedrock")
106+
@service(sdkId: "BedrockNoBearerAuth")
107+
service BedrockNoBearerAuth {
108+
version: "1.0.0"
109+
}
110+
""".trimIndent().toSmithyModel()
111+
112112
assertFalse {
113113
EnvironmentBearerTokenCustomization()
114114
.enabledForService(noBearerAuthModel, noBearerAuthModel.defaultSettings())

‎services/bedrock/e2eTest/src/BedrockEnvironmentBearerTokenTest.kt‎

Lines changed: 33 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,10 @@ import aws.smithy.kotlin.runtime.operation.ExecutionContext
2323
import aws.smithy.kotlin.runtime.time.Instant
2424
import kotlinx.coroutines.runBlocking
2525
import org.junit.jupiter.api.Assertions.assertNotNull
26-
import org.junit.jupiter.api.TestInstance
2726
import kotlin.test.Test
2827
import kotlin.test.assertEquals
2928

3029
// Environment variable AWS_BEARER_TOKEN_BEDROCK is configured with the value "bedrock-token" for this test suite.
31-
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
3230
class BedrockEnvironmentBearerTokenTest {
3331
private fun mockHttpClient(handler: (HttpRequest) -> HttpResponse): HttpClientEngine {
3432
return object : HttpClientEngineBase("test engine") {
@@ -43,57 +41,51 @@ class BedrockEnvironmentBearerTokenTest {
4341

4442
@Test
4543
fun testAuthSchemePreferenceConfigured(): Unit = runBlocking {
46-
val client = BedrockClient.fromEnvironment {
44+
BedrockClient.fromEnvironment {
4745
region = "us-west-2"
46+
}.use { client ->
47+
val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer)
48+
assertEquals(expectedAuthSchemePreference, client.config.authSchemePreference)
4849
}
49-
50-
val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer)
51-
assertEquals(expectedAuthSchemePreference, client.config.authSchemePreference)
52-
53-
client.close()
5450
}
5551

5652
@Test
5753
fun testBearerAuthSchemePromotedToFirst(): Unit = runBlocking {
58-
var client = BedrockClient.fromEnvironment {
54+
BedrockClient.fromEnvironment {
5955
region = "us-west-2"
6056
authSchemePreference = listOf(AuthSchemeId.AwsSigV4)
57+
}.use { client ->
58+
val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer, AuthSchemeId.AwsSigV4)
59+
assertEquals(expectedAuthSchemePreference, client.config.authSchemePreference)
6160
}
6261

63-
val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer, AuthSchemeId.AwsSigV4)
64-
assertEquals(expectedAuthSchemePreference, client.config.authSchemePreference)
65-
client.close()
66-
67-
client = BedrockClient.fromEnvironment {
62+
BedrockClient.fromEnvironment {
6863
region = "us-west-2"
6964
authSchemePreference = listOf(AuthSchemeId.AwsSigV4, AuthSchemeId.HttpBearer)
65+
}.use { client ->
66+
val expectedAuthSchemePreference = listOf(AuthSchemeId.HttpBearer, AuthSchemeId.AwsSigV4)
67+
assertEquals(expectedAuthSchemePreference, client.config.authSchemePreference)
7068
}
71-
72-
assertEquals(expectedAuthSchemePreference, client.config.authSchemePreference)
73-
74-
client.close()
7569
}
7670

7771
@Test
7872
fun testBearerTokenProviderConfigured(): Unit = runBlocking {
79-
val client = BedrockClient.fromEnvironment {
73+
BedrockClient.fromEnvironment {
8074
region = "us-west-2"
81-
}
82-
83-
assertNotNull(client.config.bearerTokenProvider)
84-
85-
val token = client.config.bearerTokenProvider.resolve()
86-
assertNotNull(token)
87-
assertEquals("bedrock-token", token.token)
75+
}.use { client ->
76+
assertNotNull(client.config.bearerTokenProvider)
8877

89-
client.close()
78+
val token = client.config.bearerTokenProvider.resolve()
79+
assertNotNull(token)
80+
assertEquals("bedrock-token", token.token)
81+
}
9082
}
9183

9284
@Test
9385
fun testBearerTokenProviderFunctionality(): Unit = runBlocking {
9486
var capturedAuthHeader: String? = null
9587

96-
val client = BedrockClient.fromEnvironment {
88+
BedrockClient.fromEnvironment {
9789
region = "us-west-2"
9890
httpClient = mockHttpClient { request ->
9991
// Capture the Authorization header
@@ -105,15 +97,15 @@ class BedrockEnvironmentBearerTokenTest {
10597
body = HttpBody.Empty,
10698
)
10799
}
108-
}
100+
}.use { client ->
101+
// Make an api call to capture Authorization header
102+
client.listFoundationModels()
109103

110-
// Make an api call to capture Authorization header
111-
client.listFoundationModels()
104+
assertNotNull(capturedAuthHeader)
105+
assertEquals("Bearer bedrock-token", capturedAuthHeader)
112106

113-
assertNotNull(capturedAuthHeader)
114-
assertEquals("Bearer bedrock-token", capturedAuthHeader)
115-
116-
client.close()
107+
client.close()
108+
}
117109
}
118110

119111
@Test
@@ -127,14 +119,12 @@ class BedrockEnvironmentBearerTokenTest {
127119
override val expiration: Instant? = null
128120
}
129121
}
130-
}
131-
132-
assertNotNull(client.config.bearerTokenProvider)
122+
}.use { client ->
123+
assertNotNull(client.config.bearerTokenProvider)
133124

134-
val token = client.config.bearerTokenProvider.resolve()
135-
assertNotNull(token)
136-
assertEquals("different-bedrock-token", token.token)
137-
138-
client.close()
125+
val token = client.config.bearerTokenProvider.resolve()
126+
assertNotNull(token)
127+
assertEquals("different-bedrock-token", token.token)
128+
}
139129
}
140130
}

0 commit comments

Comments
 (0)