[test] Add a test for cancellable

Sources/AWSLambdaRuntime/LambdaRuntime.swift

@@ -82,7 +82,7 @@ public final class LambdaRuntime<Handler>: Sendable where Handler: StreamingLamb
         // The handler can be non-sendable, we want to ensure we only ever have one copy of it
         let handler = try? self.handlerStorage.get()
         guard let handler else {
-            throw LambdaRuntimeError(code: .runtimeCanOnlyBeStartedOnce)
+            throw LambdaRuntimeError(code: .handlerCanOnlyBeGetOnce)
         }
 
         // are we running inside an AWS Lambda runtime environment ?

Sources/AWSLambdaRuntime/LambdaRuntimeError.swift

@@ -34,6 +34,7 @@ package struct LambdaRuntimeError: Error {
 
         case missingLambdaRuntimeAPIEnvironmentVariable
         case runtimeCanOnlyBeStartedOnce
+        case handlerCanOnlyBeGetOnce
         case invalidPort
     }
 

Tests/AWSLambdaRuntimeTests/LambdaRuntime+ServiceLifeCycle.swift

@@ -0,0 +1,46 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftAWSLambdaRuntime open source project
+//
+// Copyright (c) 2024 Apple Inc. and the SwiftAWSLambdaRuntime project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftAWSLambdaRuntime project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+#if ServiceLifecycleSupport
+@testable import AWSLambdaRuntime
+import ServiceLifecycle
+import Testing
+import Logging
+
+@Suite
+struct LambdaRuntimeServiceLifecycleTests {
+    @Test
+    func testLambdaRuntimeGracefulShutdown() async throws {
+        let runtime = LambdaRuntime {
+            (event: String, context: LambdaContext) in
+            "Hello \(event)"
+        }
+
+        let serviceGroup = ServiceGroup(
+            services: [runtime],
+            gracefulShutdownSignals: [.sigterm, .sigint],
+            logger: Logger(label: "TestLambdaRuntimeGracefulShutdown")
+        )
+        try await withThrowingTaskGroup(of: Void.self) { group in
+            group.addTask {
+                try await serviceGroup.run()
+            }
+            // wait a small amount to ensure we are waiting for continuation
+            try await Task.sleep(for: .milliseconds(100))
+
+            await serviceGroup.triggerGracefulShutdown()
+        }
+    }
+}
+#endif

Tests/AWSLambdaRuntimeTests/LambdaRuntimeClientTests.swift

@@ -15,7 +15,6 @@
 import Logging
 import NIOCore
 import NIOPosix
-import ServiceLifecycle
 import Testing
 
 import struct Foundation.UUID
@@ -90,7 +89,7 @@ struct LambdaRuntimeClientTests {
     }
 
     @Test
-    func testCancellation() async throws {
+    func testRuntimeClientCancellation() async throws {
         struct HappyBehavior: LambdaServerBehavior {
             let requestId = UUID().uuidString
             let event = "hello"
@@ -140,28 +139,4 @@ struct LambdaRuntimeClientTests {
             }
         }
     }
-    #if ServiceLifecycleSupport
-    @Test
-    func testLambdaRuntimeGracefulShutdown() async throws {
-        let runtime = LambdaRuntime {
-            (event: String, context: LambdaContext) in
-            "Hello \(event)"
-        }
-
-        let serviceGroup = ServiceGroup(
-            services: [runtime],
-            gracefulShutdownSignals: [.sigterm, .sigint],
-            logger: Logger(label: "TestLambdaRuntimeGracefulShutdown")
-        )
-        try await withThrowingTaskGroup(of: Void.self) { group in
-            group.addTask {
-                try await serviceGroup.run()
-            }
-            // wait a small amount to ensure we are waiting for continuation
-            try await Task.sleep(for: .milliseconds(100))
-
-            await serviceGroup.triggerGracefulShutdown()
-        }
-    }
-    #endif
 }

Tests/AWSLambdaRuntimeTests/LambdaRuntimeTests.swift

@@ -50,17 +50,24 @@ struct LambdaRuntimeTests {
             }
 
             // wait a small amount to ensure runtime1 task is started
-            try await Task.sleep(for: .seconds(0.5))
+            // on GH Actions, it might take a bit longer to start the runtime
+            try await Task.sleep(for: .seconds(2))
 
             // Running the second runtime should trigger LambdaRuntimeError
-            await #expect(throws: LambdaRuntimeError.self) {
-                try await runtime2.run()
+            // start the first runtime
+            taskGroup.addTask {
+                await #expect(throws: LambdaRuntimeError.self) {
+                    try await runtime2.run()
+                }
             }
 
             // cancel runtime 1 / task 1
             taskGroup.cancelAll()
         }
 
+        // wait a small amount to ensure everything is cancelled and cleanup
+        try await Task.sleep(for: .seconds(1))
+
         // Running the second runtime should work now
         try await withThrowingTaskGroup(of: Void.self) { taskGroup in
             taskGroup.addTask {
@@ -75,6 +82,53 @@ struct LambdaRuntimeTests {
             taskGroup.cancelAll()
         }
     }
+    @Test("run() must be cancellable")
+    func testLambdaRuntimeCancellable() async throws {
+
+        let logger = Logger(label: "LambdaRuntimeTests.RuntimeCancellable")
+        // create a runtime
+        let runtime = LambdaRuntime(
+            handler: MockHandler(),
+            eventLoop: Lambda.defaultEventLoop,
+            logger: logger
+        )
+
+        // Running the runtime with structured concurrency
+        // Task group returns when all tasks are completed.
+        // Even cancelled tasks must cooperatlivly complete
+        await #expect(throws: Never.self) {
+            try await withThrowingTaskGroup(of: Void.self) { taskGroup in
+                taskGroup.addTask {
+                    logger.trace("--- launching runtime ----")
+                    try await runtime.run()
+                }
+
+                // Add a timeout task to the group
+                taskGroup.addTask {
+                    logger.trace("--- launching timeout task ----")
+                    try await Task.sleep(for: .seconds(5))
+                    if Task.isCancelled { return }
+                    logger.trace("--- throwing timeout error ----")
+                    throw TestError.timeout  // Fail the test if the timeout triggers
+                }
+
+                do {
+                    // Wait for the runtime to start
+                    logger.trace("--- waiting for runtime to start ----")
+                    try await Task.sleep(for: .seconds(1))
+
+                    // Cancel all tasks, this should not throw an error
+                    // and should allow the runtime to complete gracefully
+                    logger.trace("--- cancel all tasks ----")
+                    taskGroup.cancelAll()  // Cancel all tasks
+                } catch {
+                    logger.error("--- catch an error: \(error)")
+                    throw error  // Propagate the error to fail the test
+                }
+            }
+        }
+
+    }
 }
 
 struct MockHandler: StreamingLambdaHandler {
@@ -86,3 +140,15 @@ struct MockHandler: StreamingLambdaHandler {
 
     }
 }
+
+// Define a custom error for timeout
+enum TestError: Error, CustomStringConvertible {
+    case timeout
+
+    var description: String {
+        switch self {
+        case .timeout:
+            return "Test timed out waiting for the task to complete."
+        }
+    }
+}

Tests/AWSLambdaRuntimeTests/PoolTests.swift

@@ -37,7 +37,7 @@ struct PoolTests {
     }
 
     @Test
-    func testCancellation() async throws {
+    func testPoolCancellation() async throws {
         let pool = LambdaHTTPServer.Pool<String>()
 
         // Create a task that will be cancelled