@@ -238,4 +238,54 @@ struct LambdaRuntimeClientTests {
238238 }
239239 }
240240 }
241+
242+ @Test ( " Server closing the connection when waiting for next invocation throws an error " )
243+ func testChannelCloseFutureWithWaitingForNextInvocation( ) async throws {
244+ struct DisconnectBehavior : LambdaServerBehavior {
245+ func getInvocation( ) -> GetInvocationResult {
246+ // Return "disconnect" to trigger server closing the connection
247+ . success( ( " disconnect " , " 0 " ) )
248+ }
249+
250+ func processResponse( requestId: String , response: String ? ) -> Result < Void , ProcessResponseError > {
251+ Issue . record ( " should not process response " )
252+ return . failure( . internalServerError)
253+ }
254+
255+ func processError( requestId: String , error: ErrorResponse ) -> Result < Void , ProcessErrorError > {
256+ Issue . record ( " should not report error " )
257+ return . failure( . internalServerError)
258+ }
259+
260+ func processInitError( error: ErrorResponse ) -> Result < Void , ProcessErrorError > {
261+ Issue . record ( " should not report init error " )
262+ return . failure( . internalServerError)
263+ }
264+ }
265+
266+ try await withMockServer ( behaviour: DisconnectBehavior ( ) ) { port in
267+ let configuration = LambdaRuntimeClient . Configuration ( ip: " 127.0.0.1 " , port: port)
268+
269+ try await LambdaRuntimeClient . withRuntimeClient (
270+ configuration: configuration,
271+ eventLoop: NIOSingletons . posixEventLoopGroup. next ( ) ,
272+ logger: self . logger
273+ ) { runtimeClient in
274+ do {
275+ // This should fail when server closes connection
276+ let _ = try await runtimeClient. nextInvocation ( )
277+ Issue . record ( " Expected connection error but got successful invocation " )
278+
279+ // Verify we get an error when the connection is closed.
280+ // the error is either a ChannelError or a LambdaRuntimeError
281+ } catch let error as LambdaRuntimeError {
282+ #expect( error. code == . connectionToControlPlaneLost)
283+ } catch let error as ChannelError {
284+ #expect( error == . ioOnClosedChannel)
285+ } catch {
286+ Issue . record ( " Unexpected error type: \( error) " )
287+ }
288+ }
289+ }
290+ }
241291}
0 commit comments