From 9c0398bed04f3f5b6dd6f7b118ec6ef1f625f617 Mon Sep 17 00:00:00 2001
From: Alexander Melnyk <amelnyk@amazon.com>
Date: Tue, 4 Jul 2023 12:21:18 +0200
Subject: [PATCH] add useragent to v3 sdk clients

---
 .../tracer/src/provider/ProviderService.ts    |  3 +++
 .../tracer/tests/unit/ProviderService.test.ts | 24 ++++++++++++++++---
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/packages/tracer/src/provider/ProviderService.ts b/packages/tracer/src/provider/ProviderService.ts
index b3397cbc03..380bdc7dd4 100644
--- a/packages/tracer/src/provider/ProviderService.ts
+++ b/packages/tracer/src/provider/ProviderService.ts
@@ -18,6 +18,7 @@ import {
   setLogger,
   Logger,
 } from 'aws-xray-sdk-core';
+import { addUserAgentMiddleware } from '@aws-lambda-powertools/commons';
 
 class ProviderService implements ProviderServiceInterface {
   public captureAWS<T>(awssdk: T): T {
@@ -29,6 +30,8 @@ class ProviderService implements ProviderServiceInterface {
   }
 
   public captureAWSv3Client<T>(service: T): T {
+    addUserAgentMiddleware(service, 'tracer');
+
     // Type must be aliased as any because of this https://github.com/aws/aws-xray-sdk-node/issues/439#issuecomment-859715660
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     return captureAWSv3Client(service as any);
diff --git a/packages/tracer/tests/unit/ProviderService.test.ts b/packages/tracer/tests/unit/ProviderService.test.ts
index c44b491059..6d32f510ec 100644
--- a/packages/tracer/tests/unit/ProviderService.test.ts
+++ b/packages/tracer/tests/unit/ProviderService.test.ts
@@ -6,23 +6,25 @@
 
 import { ProviderService } from '../../src/provider';
 import {
+  captureAsyncFunc,
   captureAWS,
   captureAWSClient,
   captureAWSv3Client,
-  captureAsyncFunc,
-  captureHTTPsGlobal,
   captureFunc,
+  captureHTTPsGlobal,
   getNamespace,
   getSegment,
+  Segment,
   setContextMissingStrategy,
   setDaemonAddress,
   setLogger,
   setSegment,
   Subsegment,
-  Segment,
 } from 'aws-xray-sdk-core';
 import http from 'http';
 import https from 'https';
+import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
+import * as UserAgentMiddleware from '@aws-lambda-powertools/commons/lib/userAgentMiddleware';
 
 jest.mock('aws-xray-sdk-core', () => ({
   ...jest.requireActual('aws-xray-sdk-core'),
@@ -85,6 +87,22 @@ describe('Class: ProviderService', () => {
       expect(captureAWSv3Client).toHaveBeenCalledTimes(1);
       expect(captureAWSv3Client).toHaveBeenCalledWith({});
     });
+
+    test('when called, it adds the correct user agent middleware', () => {
+      // Prepare
+      const provider: ProviderService = new ProviderService();
+
+      // spy on addUserAgentMiddleware
+      const spy = jest.spyOn(UserAgentMiddleware, 'addUserAgentMiddleware');
+
+      // Act
+      const dynamoDBClient = new DynamoDBClient({});
+      provider.captureAWSv3Client(dynamoDBClient);
+
+      // Assess
+      expect(spy).toHaveBeenCalledTimes(1);
+      expect(spy).toHaveBeenCalledWith(dynamoDBClient, 'tracer');
+    });
   });
 
   describe('Method: captureAsyncFunc', () => {