Add test for Sdk instrumentation patch

Author: blairhyy-amazon

aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts

@@ -14,6 +14,7 @@ import {
   TextMapSetter,
   INVALID_SPAN_CONTEXT,
   SpanStatusCode,
+  ROOT_CONTEXT,
 } from '@opentelemetry/api';
 import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
 import { Instrumentation } from '@opentelemetry/instrumentation';
@@ -692,6 +693,60 @@ describe('InstrumentationPatchTest', () => {
       });
     });
 
+    it('Prevents recursion when credentials provider makes STS calls', async () => {
+      // Track credentials provider calls and skip flag state
+      let credentialsProviderCallCount = 0;
+      let skipCredentialCaptureValue = false;
+
+      // Mock span for attribute capture
+      const mockSpan = { setAttribute: sinon.stub() };
+      sinon.stub(trace, 'getSpan').returns(mockSpan as unknown as Span);
+
+      // Patch ROOT_CONTEXT to capture setValue calls
+      const originalSetValue = ROOT_CONTEXT.setValue;
+      ROOT_CONTEXT.setValue = (key: symbol, value: any) => {
+        if (key.toString().includes('skip-credential-capture')) {
+          skipCredentialCaptureValue = value;
+        }
+        return ROOT_CONTEXT;
+      };
+      sinon.stub(context, 'active').returns(ROOT_CONTEXT);
+
+      // Capture middleware added to stack
+      const middlewareStack: any[] = [];
+      const recursiveSdkSend = extractAwsSdkInstrumentation(PATCHED_INSTRUMENTATIONS)
+        ['_getV3SmithyClientSendPatch'](() => Promise.resolve())
+        .bind({
+          middlewareStack: { add: (middleware: any, config: any) => middlewareStack.push([middleware, config]) },
+          config: {
+            // Credentials provider that recursively calls SDK (simulates STS)
+            credentials: async () => {
+              credentialsProviderCallCount++;
+              await recursiveSdkSend({}, null);
+              return { accessKeyId: 'test-access-key' };
+            },
+            region: () => Promise.resolve('us-west-2'),
+          },
+        });
+
+      // Initial SDK call triggers middleware setup
+      await recursiveSdkSend({}, null);
+
+      // Execute credentials extraction middleware
+      await middlewareStack[1][0](() => Promise.resolve(), null)({});
+
+      // Verify recursion prevention: only one credentials call despite recursive setup
+      expect(credentialsProviderCallCount).toBe(1);
+      expect(skipCredentialCaptureValue).toBe(true);
+      expect(
+        mockSpan.setAttribute.calledWith(AWS_ATTRIBUTE_KEYS.AWS_AUTH_ACCOUNT_ACCESS_KEY, 'test-access-key')
+      ).toBeTruthy();
+
+      // Restore original setValue
+      ROOT_CONTEXT.setValue = originalSetValue;
+      sinon.restore();
+    });
+
     it('injects trace context header into request via propagator', async () => {
       lambda = new Lambda({
         region: region,