From 4edd6c26e72686f348fb4d2512cf2a6c628c1d5f Mon Sep 17 00:00:00 2001
From: Josh Pinkney <jpink@amazon.com>
Date: Fri, 31 Jan 2025 14:33:39 -0500
Subject: [PATCH] test(core): assertTelemetry should accept partial metrics

---
 .../src/test/codewhisperer/startSecurityScan.test.ts  | 11 +++++------
 packages/core/src/test/shared/telemetry/spans.test.ts |  6 +++---
 packages/core/src/test/testUtil.ts                    |  6 +++---
 3 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/packages/core/src/test/codewhisperer/startSecurityScan.test.ts b/packages/core/src/test/codewhisperer/startSecurityScan.test.ts
index 59a517e997c..635de8603d7 100644
--- a/packages/core/src/test/codewhisperer/startSecurityScan.test.ts
+++ b/packages/core/src/test/codewhisperer/startSecurityScan.test.ts
@@ -26,7 +26,6 @@ import {
     scansLimitReachedErrorMessage,
 } from '../../codewhisperer/models/constants'
 import * as model from '../../codewhisperer/models/model'
-import { CodewhispererSecurityScan } from '../../shared/telemetry/telemetry.gen'
 import * as errors from '../../shared/errors'
 import * as timeoutUtils from '../../shared/utilities/timeoutUtils'
 import { SecurityIssuesTree } from '../../codewhisperer'
@@ -232,7 +231,7 @@ describe('startSecurityScan', function () {
             codewhispererCodeScanIssuesWithFixes: 0,
             codewhispererCodeScanScope: 'PROJECT',
             passive: false,
-        } as CodewhispererSecurityScan)
+        })
     })
 
     it('Should cancel a scan if a newer one has started', async function () {
@@ -261,7 +260,7 @@ describe('startSecurityScan', function () {
                 result: 'Cancelled',
                 reasonDesc: 'Security scan stopped by user.',
                 reason: 'DefaultError',
-            } as unknown as CodewhispererSecurityScan,
+            },
             {
                 result: 'Succeeded',
             },
@@ -323,7 +322,7 @@ describe('startSecurityScan', function () {
             reason: 'CodeScanJobFailedError',
             reasonDesc: 'CodeScanJobFailedError: Security scan failed.',
             passive: false,
-        } as unknown as CodewhispererSecurityScan)
+        })
     })
 
     it('Should show notification when throttled for project scans', async function () {
@@ -353,7 +352,7 @@ describe('startSecurityScan', function () {
             reason: 'ThrottlingException',
             reasonDesc: `ThrottlingException: Maximum com.amazon.aws.codewhisperer.StartCodeAnalysis reached for this month.`,
             passive: false,
-        } as unknown as CodewhispererSecurityScan)
+        })
     })
 
     it('Should set monthly quota exceeded when throttled for auto file scans', async function () {
@@ -385,6 +384,6 @@ describe('startSecurityScan', function () {
             reason: 'ThrottlingException',
             reasonDesc: 'ThrottlingException: Maximum file scans count reached for this month',
             passive: true,
-        } as unknown as CodewhispererSecurityScan)
+        })
     })
 })
diff --git a/packages/core/src/test/shared/telemetry/spans.test.ts b/packages/core/src/test/shared/telemetry/spans.test.ts
index 8ef782943c1..7c0daacadc1 100644
--- a/packages/core/src/test/shared/telemetry/spans.test.ts
+++ b/packages/core/src/test/shared/telemetry/spans.test.ts
@@ -102,7 +102,7 @@ describe('TelemetryTracer', function () {
             })
 
             assertTelemetry(metricName, { result: 'Succeeded', source: 'bar' })
-            assertTelemetry('apigateway_copyUrl', {} as MetricShapes['apigateway_copyUrl'])
+            assertTelemetry('apigateway_copyUrl', {})
         })
 
         it('writes to all new spans in the current context', function () {
@@ -112,7 +112,7 @@ describe('TelemetryTracer', function () {
             })
 
             assertTelemetry(metricName, { result: 'Succeeded', source: 'bar' })
-            assertTelemetry('apigateway_copyUrl', { result: 'Succeeded', source: 'bar' } as any)
+            assertTelemetry('apigateway_copyUrl', { result: 'Succeeded', source: 'bar' })
         })
 
         it('does not propagate state outside of the execution', function () {
@@ -310,7 +310,7 @@ describe('TelemetryTracer', function () {
             it('attaches the parent id to the child span', function () {
                 tracer.run(metricName, () => tracer.run(nestedName, () => {}))
                 assertTelemetry(metricName, { result: 'Succeeded' })
-                assertTelemetry(nestedName, { result: 'Succeeded', parentId: testId } as any)
+                assertTelemetry(nestedName, { result: 'Succeeded', parentId: testId })
             })
 
             it('should set trace id', function () {
diff --git a/packages/core/src/test/testUtil.ts b/packages/core/src/test/testUtil.ts
index e2ff480ea27..22b9bf6fcad 100644
--- a/packages/core/src/test/testUtil.ts
+++ b/packages/core/src/test/testUtil.ts
@@ -313,15 +313,15 @@ export function assertNoTelemetryMatch(re: RegExp | string): void | never {
  */
 export function assertTelemetry<K extends MetricName>(
     name: K,
-    expected: MetricShapes[K] | MetricShapes[K][]
+    expected: Partial<MetricShapes[K]> | Partial<MetricShapes[K]>[]
 ): void | never
 export function assertTelemetry<K extends MetricName>(
     name: K,
-    expected: MetricShapes[MetricName] | MetricShapes[MetricName][]
+    expected: Partial<MetricShapes[MetricName]> | Partial<MetricShapes[MetricName]>[]
 ): void | never
 export function assertTelemetry<K extends MetricName>(
     name: K,
-    expected: MetricShapes[K] | MetricShapes[K][]
+    expected: Partial<MetricShapes[K]> | Partial<MetricShapes[K]>[]
 ): void | never {
     const expectedList = Array.isArray(expected) ? expected : [expected]
     const query = { metricName: name }