refactor: let ScaleError convert itself to batch item failures, add tests

This lets us move some logic from `processFleetResult`, simplifying it.
We're additionally adding some bounds checking here, and testing it all.

Author: iainlane

lambdas/functions/control-plane/src/lambda.ts

@@ -62,11 +62,7 @@ export async function scaleUpHandler(event: SQSEvent, context: Context): Promise
     return { batchItemFailures };
   } catch (e) {
     if (e instanceof ScaleError) {
-      for (let i = 0; i < e.failedInstanceCount; i++) {
-        batchItemFailures.push({
-          itemIdentifier: sqsMessages[i].messageId,
-        });
-      }
+      batchItemFailures.push(...e.toBatchItemFailures(sqsMessages));
       logger.warn(`${e.detailedMessage} A retry will be attempted via SQS.`, { error: e });
     } else {
       logger.error(`Error processing batch (size: ${sqsMessages.length}): ${(e as Error).message}, ignoring batch`, {

lambdas/functions/control-plane/src/scale-runners/ScaleError.test.ts

@@ -0,0 +1,72 @@
+import { describe, expect, it } from 'vitest';
+import type { ActionRequestMessageSQS } from './scale-up';
+import ScaleError from './ScaleError';
+
+describe('ScaleError', () => {
+  describe('detailedMessage', () => {
+    it('should format message for single instance failure', () => {
+      const error = new ScaleError('Failed to create fleet', 1);
+
+      expect(error.detailedMessage).toBe('Failed to create fleet (Failed to create 1 instance)');
+    });
+
+    it('should format message for multiple instance failures', () => {
+      const error = new ScaleError('Failed to create fleet', 3);
+
+      expect(error.detailedMessage).toBe('Failed to create fleet (Failed to create 3 instances)');
+    });
+  });
+
+  describe('toBatchItemFailures', () => {
+    const mockMessages: ActionRequestMessageSQS[] = [
+      { messageId: 'msg-1', id: 1, eventType: 'workflow_job' },
+      { messageId: 'msg-2', id: 2, eventType: 'workflow_job' },
+      { messageId: 'msg-3', id: 3, eventType: 'workflow_job' },
+      { messageId: 'msg-4', id: 4, eventType: 'workflow_job' },
+    ];
+
+    it.each([
+      { failedCount: 1, expected: [{ itemIdentifier: 'msg-1' }], description: 'default instance count' },
+      {
+        failedCount: 2,
+        expected: [{ itemIdentifier: 'msg-1' }, { itemIdentifier: 'msg-2' }],
+        description: 'less than message count',
+      },
+      {
+        failedCount: 4,
+        expected: [
+          { itemIdentifier: 'msg-1' },
+          { itemIdentifier: 'msg-2' },
+          { itemIdentifier: 'msg-3' },
+          { itemIdentifier: 'msg-4' },
+        ],
+        description: 'equal to message count',
+      },
+      {
+        failedCount: 10,
+        expected: [
+          { itemIdentifier: 'msg-1' },
+          { itemIdentifier: 'msg-2' },
+          { itemIdentifier: 'msg-3' },
+          { itemIdentifier: 'msg-4' },
+        ],
+        description: 'more than message count',
+      },
+      { failedCount: 0, expected: [], description: 'zero failed instances' },
+      { failedCount: -1, expected: [], description: 'negative failed instances' },
+      { failedCount: -10, expected: [], description: 'large negative failed instances' },
+    ])('should handle $description (failedCount=$failedCount)', ({ failedCount, expected }) => {
+      const error = new ScaleError('Test error', failedCount);
+      const failures = error.toBatchItemFailures(mockMessages);
+
+      expect(failures).toEqual(expected);
+    });
+
+    it('should handle empty message array', () => {
+      const error = new ScaleError('Test error', 3);
+      const failures = error.toBatchItemFailures([]);
+
+      expect(failures).toEqual([]);
+    });
+  });
+});

lambdas/functions/control-plane/src/scale-runners/ScaleError.ts

@@ -1,3 +1,6 @@
+import type { SQSBatchItemFailure } from 'aws-lambda';
+import type { ActionRequestMessageSQS } from './scale-up';
+
 class ScaleError extends Error {
   constructor(
     message: string,
@@ -13,6 +16,17 @@ class ScaleError extends Error {
   public get detailedMessage(): string {
     return `${this.message} (Failed to create ${this.failedInstanceCount} instance${this.failedInstanceCount !== 1 ? 's' : ''})`;
   }
+
+  /**
+   * Generate SQS batch item failures for the failed instances
+   */
+  public toBatchItemFailures(messages: ActionRequestMessageSQS[]): SQSBatchItemFailure[] {
+    // Ensure we don't retry negative counts or more messages than available
+    const messagesToRetry = Math.max(0, Math.min(this.failedInstanceCount, messages.length));
+    return messages.slice(0, messagesToRetry).map(({ messageId }) => ({
+      itemIdentifier: messageId,
+    }));
+  }
 }
 
-export default ScaleError;
+export default ScaleError;