test: invoke twice back-to-back with a single indexing wait

Replaces the two-window/two-5-min-indexing-wait design with one window and
one wait, comparing tagged vs. total invocation counts (new
getMetricCountWithTag helper) instead of checking existence per window.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: lym953

integration-tests/tests/durable-cold-start.test.ts

@@ -1,6 +1,6 @@
-// Verifies the invocations metric has durable_function:true only for a durable function, on
-// both the cold-start invocation and a subsequent warm invocation.
-import { hasMetricWithTag } from './utils/datadog';
+// Verifies the invocations metric has durable_function:true only for a durable function, and
+// that the tag persists across a second (warm) invocation instead of only being set on cold start.
+import { getMetricCount, getMetricCountWithTag } from './utils/datadog';
 import { forceColdStart, invokeLambda } from './utils/lambda';
 import { IDENTIFIER, DEFAULT_DATADOG_INDEXING_WAIT_MS } from '../config';
 
@@ -9,10 +9,8 @@ const stackName = `${IDENTIFIER}-durable-cold-start`;
 const INVOCATIONS_METRIC = 'aws.lambda.enhanced.invocations';
 
 describe('Durable Function Cold-Start Metric Tag Integration Tests', () => {
-  let coldStartWindowStart: number;
-  let coldStartWindowEnd: number;
-  let warmInvokeWindowStart: number;
-  let warmInvokeWindowEnd: number;
+  let windowStart: number;
+  let windowEnd: number;
 
   const durableFunctionName = `${stackName}-durable-lambda`;
   const nonDurableFunctionName = `${stackName}-non-durable-lambda`;
@@ -24,70 +22,60 @@ describe('Durable Function Cold-Start Metric Tag Integration Tests', () => {
       invokeLambda(nonDurableFunctionName),
     ]);
 
-  const waitForIndexing = () =>
-    new Promise((resolve) =>
-      setTimeout(resolve, DEFAULT_DATADOG_INDEXING_WAIT_MS),
-    );
-
   beforeAll(async () => {
     const functionNames = [durableFunctionName, nonDurableFunctionName];
     await Promise.all(functionNames.map((fn) => forceColdStart(fn)));
 
     // Back up 60s so the metric's rollup bucket falls inside the query range.
-    coldStartWindowStart = Date.now() - 60_000;
-    await invokeBoth(); // invocation #1 (cold start)
-    await waitForIndexing();
-    coldStartWindowEnd = Date.now();
+    windowStart = Date.now() - 60_000;
+
+    // Invoke twice back-to-back (cold start, then warm) and wait for indexing once, rather than
+    // waiting after each invocation. Whether the tag was set on both invocations (rather than
+    // only the first) is checked below by comparing tagged vs. total invocation counts.
+    await invokeBoth();
+    await invokeBoth();
 
-    warmInvokeWindowStart = Date.now() - 60_000;
-    await invokeBoth(); // invocation #2 (warm)
-    await waitForIndexing();
-    warmInvokeWindowEnd = Date.now();
+    await new Promise((resolve) =>
+      setTimeout(resolve, DEFAULT_DATADOG_INDEXING_WAIT_MS),
+    );
+    windowEnd = Date.now();
 
-    console.log('Lambdas invoked twice and indexing waits complete');
-  }, 3600000);
+    console.log('Lambdas invoked twice and indexing wait complete');
+  }, 1800000);
 
-  it('durable function cold-start invocation metric should have durable_function:true', async () => {
-    const hasTag = await hasMetricWithTag(
+  it('durable function invocation metric should have durable_function:true on every invocation', async () => {
+    const totalCount = await getMetricCount(
       INVOCATIONS_METRIC,
       durableFunctionName,
-      'durable_function:true',
-      coldStartWindowStart,
-      coldStartWindowEnd,
+      windowStart,
+      windowEnd,
     );
-    expect(hasTag).toBe(true);
-  });
-
-  it('non-durable function cold-start invocation metric should NOT have durable_function:true', async () => {
-    const hasTag = await hasMetricWithTag(
+    const taggedCount = await getMetricCountWithTag(
       INVOCATIONS_METRIC,
-      nonDurableFunctionName,
+      durableFunctionName,
       'durable_function:true',
-      coldStartWindowStart,
-      coldStartWindowEnd,
+      windowStart,
+      windowEnd,
     );
-    expect(hasTag).toBe(false);
+    expect(totalCount).toBeGreaterThanOrEqual(2);
+    expect(taggedCount).toBe(totalCount);
   });
 
-  it('durable function warm invocation metric should still have durable_function:true', async () => {
-    const hasTag = await hasMetricWithTag(
+  it('non-durable function invocation metric should NOT have durable_function:true on any invocation', async () => {
+    const totalCount = await getMetricCount(
       INVOCATIONS_METRIC,
-      durableFunctionName,
-      'durable_function:true',
-      warmInvokeWindowStart,
-      warmInvokeWindowEnd,
+      nonDurableFunctionName,
+      windowStart,
+      windowEnd,
     );
-    expect(hasTag).toBe(true);
-  });
-
-  it('non-durable function warm invocation metric should still NOT have durable_function:true', async () => {
-    const hasTag = await hasMetricWithTag(
+    const taggedCount = await getMetricCountWithTag(
       INVOCATIONS_METRIC,
       nonDurableFunctionName,
       'durable_function:true',
-      warmInvokeWindowStart,
-      warmInvokeWindowEnd,
+      windowStart,
+      windowEnd,
     );
-    expect(hasTag).toBe(false);
+    expect(totalCount).toBeGreaterThanOrEqual(2);
+    expect(taggedCount).toBe(0);
   });
 });

integration-tests/tests/utils/datadog.ts

@@ -370,6 +370,40 @@ export async function getMetricCount(
   return pointlist.reduce((acc, [, value]) => acc + (value || 0), 0);
 }
 
+/**
+ * Same as `getMetricCount`, but filtered by an additional tag (e.g. "durable_function:true").
+ * Used to check whether every emission of a metric over a window carried a given tag, by
+ * comparing against the untagged `getMetricCount` total for the same window.
+ */
+export async function getMetricCountWithTag(
+  metricName: string,
+  functionName: string,
+  tagFilter: string,
+  fromTime: number,
+  toTime: number,
+): Promise<number> {
+  const baseFunctionName = getServiceName(functionName).toLowerCase();
+  const query = `sum:${metricName}{functionname:${baseFunctionName},${tagFilter}}.as_count()`;
+
+  console.log(`Querying metric count with tag filter: ${query}`);
+
+  const response = await requestWithRetry(() => datadogClient.get('/api/v1/query', {
+    params: {
+      query,
+      from: Math.floor(fromTime / 1000),
+      to: Math.floor(toTime / 1000),
+    },
+  }), query);
+
+  const series = response.data.series || [];
+  if (series.length === 0) {
+    return 0;
+  }
+
+  const pointlist: [number, number][] = series[0].pointlist || [];
+  return pointlist.reduce((acc, [, value]) => acc + (value || 0), 0);
+}
+
 async function getMetrics(
   metricName: string,
   functionName: string,