[test] Add a test for cancellable (#529)

- Add a test on `LambdaRuntime` for cancellable.
- Move the service life cycle test to its own file

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

@@ -41,40 +41,79 @@ struct LambdaRuntimeTests {
         )
 
         try await withThrowingTaskGroup(of: Void.self) { taskGroup in
+
             // start the first runtime
             taskGroup.addTask {
-                // ChannelError will be thrown when we cancel the task group
-                await #expect(throws: ChannelError.self) {
-                    try await runtime1.run()
-                }
+                // will throw LambdaRuntimeError when run() is called second or ChannelError when cancelled
+                try await runtime1.run()
             }
 
             // wait a small amount to ensure runtime1 task is started
             try await Task.sleep(for: .seconds(0.5))
 
-            // Running the second runtime should trigger LambdaRuntimeError
-            await #expect(throws: LambdaRuntimeError.self) {
+            // start the second runtime
+            taskGroup.addTask {
+                // will throw LambdaRuntimeError when run() is called second or ChannelError when cancelled
                 try await runtime2.run()
             }
 
-            // cancel runtime 1 / task 1
-            taskGroup.cancelAll()
-        }
-
-        // Running the second runtime should work now
-        try await withThrowingTaskGroup(of: Void.self) { taskGroup in
-            taskGroup.addTask {
-                // ChannelError will be thrown when we cancel the task group
-                await #expect(throws: ChannelError.self) {
-                    try await runtime2.run()
-                }
+            // get the first result (should throw a LambdaRuntimeError)
+            try await #require(throws: LambdaRuntimeError.self) {
+                try await taskGroup.next()
             }
 
-            // Set timeout and cancel the runtime 2
-            try await Task.sleep(for: .seconds(1))
+            // cancel the group to end the test
             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 +125,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