Refactor + add test for reportError

Author: aryan-25

Tests/AWSLambdaRuntimeCoreTests/LambdaMockClient.swift

@@ -28,23 +28,27 @@ struct LambdaMockWriter: LambdaResponseStreamWriter {
         try await self.underlying.write(buffer)
     }
 
-    consuming func finish() async throws {
+    func finish() async throws {
         try await self.underlying.finish()
     }
 
-    consuming func writeAndFinish(_ buffer: ByteBuffer) async throws {
-        try await self.write(buffer)
-        try await self.finish()
+    func writeAndFinish(_ buffer: ByteBuffer) async throws {
+        try await self.underlying.write(buffer)
+        try await self.underlying.finish()
     }
 
     func reportError(_ error: any Error) async throws {
+        try await self.underlying.write(ByteBuffer(string: "\(error)"))
+        try await self.underlying.finish()
     }
 }
 
 enum LambdaError: Error, Equatable {
     case cannotCallNextEndpointWhenAlreadyWaitingForEvent
     case cannotCallNextEndpointWhenAlreadyProcessingAnEvent
     case cannotReportResultWhenNoEventHasBeenProcessed
+    case cancelError
+    case handlerError
 }
 
 final actor LambdaMockClient: LambdaRuntimeClientProtocol {
@@ -89,6 +93,12 @@ final actor LambdaMockClient: LambdaRuntimeClientProtocol {
             case fail(LambdaError)
         }
 
+        enum CancelNextAction {
+            case none
+
+            case cancelContinuation(CheckedContinuation<Invocation, any Error>)
+        }
+
         enum ResultAction {
             case readyForMore
 
@@ -165,6 +175,16 @@ final actor LambdaMockClient: LambdaRuntimeClientProtocol {
                 throw LambdaError.cannotReportResultWhenNoEventHasBeenProcessed
             }
         }
+
+        mutating func cancelNext() -> CancelNextAction {
+            switch self.state {
+            case .initialState, .handlerIsProcessing:
+                return .none
+            case .waitingForNextEvent(let eventArrivedHandler):
+                self.state = .initialState
+                return .cancelContinuation(eventArrivedHandler)
+            }
+        }
     }
 
     private var stateMachine: StateMachine = .init()
@@ -208,18 +228,32 @@ final actor LambdaMockClient: LambdaRuntimeClientProtocol {
     }
 
     func nextInvocation() async throws -> (Invocation, Writer) {
-        let invocation = try await withCheckedThrowingContinuation { eventArrivedHandler in
-            switch self.stateMachine.next(eventArrivedHandler) {
-            case .readyToProcess(let event):
-                eventArrivedHandler.resume(returning: event)
-            case .fail(let error):
-                eventArrivedHandler.resume(throwing: error)
-            case .wait:
-                break
+        try await withTaskCancellationHandler {
+            let invocation = try await withCheckedThrowingContinuation { eventArrivedHandler in
+                switch self.stateMachine.next(eventArrivedHandler) {
+                case .readyToProcess(let event):
+                    eventArrivedHandler.resume(returning: event)
+                case .fail(let error):
+                    eventArrivedHandler.resume(throwing: error)
+                case .wait:
+                    break
+                }
+            }
+            return (invocation, Writer(underlying: self))
+        } onCancel: {
+            Task {
+                await self.cancelNextInvocation()
             }
         }
+    }
 
-        return (invocation, Writer(underlying: self))
+    private func cancelNextInvocation() {
+        switch self.stateMachine.cancelNext() {
+        case .none:
+            break
+        case .cancelContinuation(let continuation):
+            continuation.resume(throwing: LambdaError.cancelError)
+        }
     }
 
     func write(_ buffer: ByteBuffer) async throws {

Tests/AWSLambdaRuntimeCoreTests/LambdaRunLoopTests.swift

@@ -31,23 +31,55 @@ struct LambdaRunLoopTests {
         }
     }
 
+    struct FailingHandler: StreamingLambdaHandler {
+        func handle(
+            _ event: ByteBuffer,
+            responseWriter: some LambdaResponseStreamWriter,
+            context: NewLambdaContext
+        ) async throws {
+            throw LambdaError.handlerError
+        }
+    }
+
     let mockClient = LambdaMockClient()
     let mockEchoHandler = MockEchoHandler()
+    let failingHandler = FailingHandler()
 
     @Test func testRunLoop() async throws {
-        let runLoopTask = Task { () in
-            try await Lambda.runLoop(
-                runtimeClient: self.mockClient,
-                handler: self.mockEchoHandler,
-                logger: Logger(label: "RunLoopTest")
-            )
+        let inputEvent = ByteBuffer(string: "Test Invocation Event")
+
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            group.addTask {
+                try await Lambda.runLoop(
+                    runtimeClient: self.mockClient,
+                    handler: self.mockEchoHandler,
+                    logger: Logger(label: "RunLoopTest")
+                )
+            }
+
+            let response = try await self.mockClient.invoke(event: inputEvent)
+            #expect(response == inputEvent)
+
+            group.cancelAll()
         }
+    }
 
+    @Test func testRunLoopError() async throws {
         let inputEvent = ByteBuffer(string: "Test Invocation Event")
-        let response = try await self.mockClient.invoke(event: inputEvent)
 
-        runLoopTask.cancel()
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            group.addTask {
+                try await Lambda.runLoop(
+                    runtimeClient: self.mockClient,
+                    handler: self.failingHandler,
+                    logger: Logger(label: "RunLoopTest")
+                )
+            }
+
+            let response = try await self.mockClient.invoke(event: inputEvent)
+            #expect(String(buffer: response) == "\(LambdaError.handlerError)")
 
-        #expect(response == inputEvent)
+            group.cancelAll()
+        }
     }
 }