withTelemetryContextFactory

Signed-off-by: nkomonen-amazon <[email protected]>

Author: nkomonen-amazon

packages/core/src/shared/telemetry/util.ts

@@ -426,3 +426,11 @@ export function withTelemetryContext(opts: TelemetryContextArgs) {
         })
     }
 }
+
+/**
+ * Returns a custom {@link withTelemetryContext} decorator, but with some values predefined
+ * to deduplicate boilerplate for a class that will use this decorator multiple times.
+ */
+export function withTelemetryContextFactory(predefinedOpts: Omit<TelemetryContextArgs, 'name'>) {
+    return (opts: { name: string }) => withTelemetryContext({ ...opts, ...predefinedOpts })
+}

packages/core/src/test/shared/telemetry/spans.test.ts

@@ -11,7 +11,7 @@ import { assertTelemetry, getMetrics, installFakeClock } from '../../testUtil'
 import { selectFrom } from '../../../shared/utilities/tsUtils'
 import { getAwsServiceError } from '../errors.test'
 import { sleep } from '../../../shared'
-import { withTelemetryContext } from '../../../shared/telemetry/util'
+import { withTelemetryContext, withTelemetryContextFactory } from '../../../shared/telemetry/util'
 import { SinonSandbox } from 'sinon'
 import sinon from 'sinon'
 import * as crypto from '../../../shared/crypto'
@@ -685,6 +685,27 @@ describe('TelemetryTracer', function () {
                 (e) => e === arbitraryError
             )
         })
+
+        it('Uses a custom telemetry context', function () {
+            const customWithTelemetryContext = withTelemetryContextFactory({ class: 'TestCustomContext', emit: true })
+
+            class TestCustomContext {
+                @customWithTelemetryContext({ name: 'testMethod' })
+                testMethod() {
+                    return
+                }
+            }
+
+            const inst = new TestCustomContext()
+            inst.testMethod()
+            assertTelemetry('function_call', [
+                {
+                    functionName: 'testMethod',
+                    className: 'TestCustomContext',
+                    source: 'TestCustomContext#testMethod',
+                },
+            ])
+        })
     })
 
     describe('withTraceId()', function () {