Fix reportError

Author: aryan-25

Tests/AWSLambdaRuntimeCoreTests/LambdaMockClient.swift

@@ -38,8 +38,7 @@ struct LambdaMockWriter: LambdaResponseStreamWriter {
     }
 
     func reportError(_ error: any Error) async throws {
-        try await self.underlying.write(ByteBuffer(string: "\(error)"))
-        try await self.underlying.finish()
+        await self.underlying.reportError(error)
     }
 }
 
@@ -105,6 +104,12 @@ final actor LambdaMockClient: LambdaRuntimeClientProtocol {
             case fail(LambdaError)
         }
 
+        enum CancelProcessingAction {
+            case none
+
+            case cancelContinuation(CheckedContinuation<ByteBuffer, any Error>)
+        }
+
         mutating func next(_ eventArrivedHandler: CheckedContinuation<Invocation, any Error>) -> NextAction {
             switch self.state {
             case .initialState:
@@ -185,6 +190,15 @@ final actor LambdaMockClient: LambdaRuntimeClientProtocol {
                 return .cancelContinuation(eventArrivedHandler)
             }
         }
+
+        mutating func cancelProcessing() -> CancelProcessingAction {
+            switch self.state {
+            case .initialState, .waitingForNextEvent:
+                return .none
+            case .handlerIsProcessing(_, let eventProcessedHandler):
+                return .cancelContinuation(eventProcessedHandler)
+            }
+        }
     }
 
     private var stateMachine: StateMachine = .init()
@@ -268,4 +282,13 @@ final actor LambdaMockClient: LambdaRuntimeClientProtocol {
     func finish() async throws {
         try self.stateMachine.finish()
     }
+
+    func reportError(_ error: any Error) {
+        switch self.stateMachine.cancelProcessing() {
+        case .none:
+            break
+        case .cancelContinuation(let continuation):
+            continuation.resume(throwing: error)
+        }
+    }
 }

Tests/AWSLambdaRuntimeCoreTests/LambdaRunLoopTests.swift

@@ -67,7 +67,7 @@ struct LambdaRunLoopTests {
     @Test func testRunLoopError() async throws {
         let inputEvent = ByteBuffer(string: "Test Invocation Event")
 
-        try await withThrowingTaskGroup(of: Void.self) { group in
+        await withThrowingTaskGroup(of: Void.self) { group in
             group.addTask {
                 try await Lambda.runLoop(
                     runtimeClient: self.mockClient,
@@ -76,8 +76,12 @@ struct LambdaRunLoopTests {
                 )
             }
 
-            let response = try await self.mockClient.invoke(event: inputEvent)
-            #expect(String(buffer: response) == "\(LambdaError.handlerError)")
+            await #expect(
+                throws: LambdaError.handlerError,
+                performing: {
+                    try await self.mockClient.invoke(event: inputEvent)
+                }
+            )
 
             group.cancelAll()
         }