test(clients): user agent is set

Author: justinmk3

packages/core/src/auth/sso/clients.ts

@@ -39,7 +39,7 @@ import { getUserAgent, withTelemetryContext } from '../../shared/telemetry/util'
 
 export class OidcClient {
     public constructor(
-        private readonly client: SSOOIDC,
+        public readonly client: SSOOIDC,
         private readonly clock: { Date: typeof Date }
     ) {}
 
@@ -157,7 +157,7 @@ export class SsoClient {
     }
 
     public constructor(
-        private readonly client: PromisifyClient<SSO>,
+        public readonly client: PromisifyClient<SSO>,
         private readonly provider: SsoAccessTokenProvider
     ) {}
 

packages/core/src/shared/clients/codewhispererChatClient.ts

@@ -9,14 +9,14 @@ import { AuthUtil } from '../../codewhisperer/util/authUtil'
 import { getUserAgent } from '../telemetry/util'
 
 // Create a client for featureDev streaming based off of aws sdk v3
-export async function createCodeWhispererChatStreamingClient(): Promise<CodeWhispererStreaming> {
-    const bearerToken = await AuthUtil.instance.getBearerToken()
+export async function createCodeWhispererChatStreamingClient(bearerToken?: string): Promise<CodeWhispererStreaming> {
+    bearerToken = bearerToken ?? (await AuthUtil.instance.getBearerToken())
     const cwsprConfig = getCodewhispererConfig()
     const streamingClient = new CodeWhispererStreaming({
         region: cwsprConfig.region,
         endpoint: cwsprConfig.endpoint,
         token: { token: bearerToken },
-        customUserAgent: getUserAgent(),
+        customUserAgent: getUserAgent({ includePlatform: true, includeClientId: true }),
         // SETTING max attempts to 0 FOR BETA. RE-ENABLE FOR RE-INVENT
         // Implement exponential back off starting with a base of 500ms (500 + attempt^10)
         retryStrategy: new ConfiguredRetryStrategy(0, (attempt: number) => 500 + attempt ** 10),

packages/core/src/test/shared/defaultAwsClientBuilder.test.ts

@@ -8,19 +8,37 @@ import { AWSError, Request, Service } from 'aws-sdk'
 import { version } from 'vscode'
 import { AWSClientBuilder, DefaultAWSClientBuilder } from '../../shared/awsClientBuilder'
 import { DevSettings } from '../../shared/settings'
-import { getClientId } from '../../shared/telemetry/util'
+import { getClientId, getUserAgent } from '../../shared/telemetry/util'
 import { FakeMemento } from '../fakeExtensionContext'
 import { FakeAwsContext } from '../utilities/fakeAwsContext'
 import { TestSettings } from '../utilities/testSettingsConfiguration'
 import { GlobalState } from '../../shared/globalState'
+import { createCodeWhispererChatStreamingClient } from '../../shared/clients/codewhispererChatClient'
+import { OidcClient, SsoClient } from '../../auth/sso/clients'
 
-describe('DefaultAwsClientBuilder', function () {
+describe('AwsClientBuilder', function () {
     let builder: AWSClientBuilder
 
     beforeEach(function () {
         builder = new DefaultAWSClientBuilder(new FakeAwsContext())
     })
 
+    it('service clients set user agent', async function () {
+        const userAgent = getUserAgent({ includePlatform: true, includeClientId: true })
+
+        const cwclient = await createCodeWhispererChatStreamingClient('fake-token')
+        assert.deepStrictEqual(cwclient.config.customUserAgent, [[userAgent]])
+
+        const oidcClient = OidcClient.create('us-east-2')
+        assert.deepStrictEqual(oidcClient.client.config.customUserAgent, [[userAgent]])
+
+        const ssoClient = SsoClient.create('us-east-2', {} as any)
+        assert.deepStrictEqual(ssoClient.client.config.customUserAgent, [[userAgent]])
+
+        const awsService = await builder.createAwsService(Service)
+        assert.deepStrictEqual(awsService.config.customUserAgent, userAgent)
+    })
+
     describe('createAndConfigureSdkClient', function () {
         it('includes Toolkit user-agent if no options are specified', async function () {
             const service = await builder.createAwsService(Service)