make sure connectionToControlPlaneLost error is triggered by the test

Author: sebsto

Tests/AWSLambdaRuntimeTests/LambdaRuntimeClientTests.swift

@@ -42,10 +42,10 @@ struct LambdaRuntimeClientTests {
                 .success((self.requestId, self.event))
             }
 
-            func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
+            func processResponse(requestId: String, response: String?) -> Result<String?, ProcessResponseError> {
                 #expect(self.requestId == requestId)
                 #expect(self.event == response)
-                return .success(())
+                return .success(nil)
             }
 
             func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
@@ -102,9 +102,9 @@ struct LambdaRuntimeClientTests {
             .success((self.requestId, self.event))
         }
 
-        func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
+        func processResponse(requestId: String, response: String?) -> Result<String?, ProcessResponseError> {
             #expect(self.requestId == requestId)
-            return .success(())
+            return .success(nil)
         }
 
         mutating func captureHeaders(_ headers: HTTPHeaders) {
@@ -197,10 +197,10 @@ struct LambdaRuntimeClientTests {
                 .success((self.requestId, self.event))
             }
 
-            func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
+            func processResponse(requestId: String, response: String?) -> Result<String?, ProcessResponseError> {
                 #expect(self.requestId == requestId)
                 #expect(self.event == response)
-                return .success(())
+                return .success(nil)
             }
 
             func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
@@ -239,31 +239,56 @@ struct LambdaRuntimeClientTests {
         }
     }
 
-    @Test("Server closing the connection when waiting for next invocation throws an error")
-    func testChannelCloseFutureWithWaitingForNextInvocation() async throws {
-        struct DisconnectBehavior: LambdaServerBehavior {
-            func getInvocation() -> GetInvocationResult {
-                // Return "disconnect" to trigger server closing the connection
-                .success(("disconnect", "0"))
-            }
+    struct DisconnectAfterSendingResponseBehavior: LambdaServerBehavior {
+        func getInvocation() -> GetInvocationResult {
+            .success((UUID().uuidString, "hello"))
+        }
 
-            func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError> {
-                Issue.record("should not process response")
-                return .failure(.internalServerError)
-            }
+        func processResponse(requestId: String, response: String?) -> Result<String?, ProcessResponseError> {
+            // Return "disconnect" to trigger server closing the connection
+            // after having accepted a response
+            .success("delayed-disconnect")
+        }
 
-            func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
-                Issue.record("should not report error")
-                return .failure(.internalServerError)
-            }
+        func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
+            Issue.record("should not report error")
+            return .failure(.internalServerError)
+        }
 
-            func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError> {
-                Issue.record("should not report init error")
-                return .failure(.internalServerError)
-            }
+        func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError> {
+            Issue.record("should not report init error")
+            return .failure(.internalServerError)
+        }
+    }
+
+    struct DisconnectBehavior: LambdaServerBehavior {
+        func getInvocation() -> GetInvocationResult {
+            // Return "disconnect" to trigger server closing the connection
+            .success(("disconnect", "0"))
         }
 
-        try await withMockServer(behaviour: DisconnectBehavior()) { port in
+        func processResponse(requestId: String, response: String?) -> Result<String?, ProcessResponseError> {
+            Issue.record("should not process response")
+            return .failure(.internalServerError)
+        }
+
+        func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError> {
+            Issue.record("should not report error")
+            return .failure(.internalServerError)
+        }
+
+        func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError> {
+            Issue.record("should not report init error")
+            return .failure(.internalServerError)
+        }
+    }
+
+    @Test(
+        "Server closing the connection when waiting for next invocation throws an error",
+        arguments: [DisconnectAfterSendingResponseBehavior(), DisconnectBehavior()] as [any LambdaServerBehavior]
+    )
+    func testChannelCloseFutureWithWaitingForNextInvocation(behavior: LambdaServerBehavior) async throws {
+        try await withMockServer(behaviour: behavior) { port in
             let configuration = LambdaRuntimeClient.Configuration(ip: "127.0.0.1", port: port)
 
             try await LambdaRuntimeClient.withRuntimeClient(
@@ -273,7 +298,12 @@ struct LambdaRuntimeClientTests {
             ) { runtimeClient in
                 do {
                     // This should fail when server closes connection
+                    let (_, writer) = try await runtimeClient.nextInvocation()
+                    let response = ByteBuffer(string: "hello")
+                    try await writer.writeAndFinish(response)
+
                     let _ = try await runtimeClient.nextInvocation()
+
                     Issue.record("Expected connection error but got successful invocation")
 
                 } catch let error as LambdaRuntimeError {

Tests/AWSLambdaRuntimeTests/MockLambdaServer.swift

@@ -160,6 +160,7 @@ final class HTTPHandler: ChannelInboundHandler {
         var responseStatus: HTTPResponseStatus
         var responseBody: String?
         var responseHeaders: [(String, String)]?
+        var disconnectAfterSend = false
 
         // Handle post-init-error first to avoid matching the less specific post-error suffix.
         if request.head.uri.hasSuffix(Consts.postInitErrorURL) {
@@ -202,8 +203,11 @@ final class HTTPHandler: ChannelInboundHandler {
             behavior.captureHeaders(request.head.headers)
 
             switch behavior.processResponse(requestId: String(requestId), response: requestBody) {
-            case .success:
+            case .success(let next):
                 responseStatus = .accepted
+                if next == "delayed-disconnect" {
+                    disconnectAfterSend = true
+                }
             case .failure(let error):
                 responseStatus = .init(statusCode: error.rawValue)
             }
@@ -223,14 +227,21 @@ final class HTTPHandler: ChannelInboundHandler {
         } else {
             responseStatus = .notFound
         }
-        self.writeResponse(context: context, status: responseStatus, headers: responseHeaders, body: responseBody)
+        self.writeResponse(
+            context: context,
+            status: responseStatus,
+            headers: responseHeaders,
+            body: responseBody,
+            closeConnection: disconnectAfterSend
+        )
     }
 
     func writeResponse(
         context: ChannelHandlerContext,
         status: HTTPResponseStatus,
         headers: [(String, String)]? = nil,
-        body: String? = nil
+        body: String? = nil,
+        closeConnection: Bool = false
     ) {
         var headers = HTTPHeaders(headers ?? [])
         headers.add(name: "Content-Length", value: "\(body?.utf8.count ?? 0)")
@@ -253,14 +264,19 @@ final class HTTPHandler: ChannelInboundHandler {
         }
 
         let loopBoundContext = NIOLoopBound(context, eventLoop: context.eventLoop)
-
         let keepAlive = self.keepAlive
         context.writeAndFlush(wrapOutboundOut(.end(nil))).whenComplete { result in
+            let context = loopBoundContext.value
+            if closeConnection {
+                context.close(promise: nil)
+                return
+            }
+
             if case .failure(let error) = result {
                 logger.error("write error \(error)")
             }
+
             if !keepAlive {
-                let context = loopBoundContext.value
                 context.close().whenFailure { error in
                     logger.error("close error \(error)")
                 }
@@ -271,7 +287,7 @@ final class HTTPHandler: ChannelInboundHandler {
 
 protocol LambdaServerBehavior: Sendable {
     func getInvocation() -> GetInvocationResult
-    func processResponse(requestId: String, response: String?) -> Result<Void, ProcessResponseError>
+    func processResponse(requestId: String, response: String?) -> Result<String?, ProcessResponseError>
     func processError(requestId: String, error: ErrorResponse) -> Result<Void, ProcessErrorError>
     func processInitError(error: ErrorResponse) -> Result<Void, ProcessErrorError>