Fix Unit Test and Timeout

Author: Andrea Scuderi

Sources/BreezeDynamoDBService/BreezeDynamoDBService.swift

@@ -33,7 +33,25 @@ public struct BreezeDynamoDBService: BreezeDynamoDBServing {
     private let logger: Logger
     private let awsClient: AWSClient
     private let httpClient: HTTPClient
-    
+    private let shutdownState: ShutdownState
+
+    /// Error types for BreezeDynamoDBService
+    enum BreezeDynamoDBServiceError: Error {
+        case alreadyShutdown
+    }
+
+    /// Actor to manage shutdown state safely
+    private actor ShutdownState {
+        private var isShutdown = false
+        
+        func markShutdown() throws {
+            guard !isShutdown else {
+                throw BreezeDynamoDBServiceError.alreadyShutdown
+            }
+            isShutdown = true
+        }
+    }
+        
     /// Initializes the BreezeDynamoDBService with the provided configuration.
     /// - Parameters:
     ///   - config: The configuration for the DynamoDB service.
@@ -76,6 +94,7 @@ public struct BreezeDynamoDBService: BreezeDynamoDBServing {
             tableName: config.tableName,
             keyName: config.keyName
         )
+        self.shutdownState = ShutdownState()
         logger.info("DBManager is ready.")
     }
 
@@ -89,9 +108,10 @@ public struct BreezeDynamoDBService: BreezeDynamoDBServing {
     /// - Throws: An error if the shutdown process fails.
     /// This method ensures that the AWS client and HTTP client are properly shutdown before marking the service as shutdown.
     /// It also logs the shutdown process.
-    /// This method is idempotent;
-    /// - Important: This method must be called at leat once to ensure that resources are released properly. If the method is not called, it will lead to a crash.
+    /// This method is idempotent and will throw if called multiple times to prevent double shutdown.
+    /// - Important: This method must be called at least once to ensure that resources are released properly. If the method is not called, it will lead to a crash.
     public func onGracefulShutdown() async throws {
+        try await shutdownState.markShutdown()
         logger.info("Stopping DynamoDBService...")
         try await awsClient.shutdown()
         logger.info("DynamoDBService is stopped.")

Sources/BreezeLambdaAPI/BreezeLambdaAPI.swift

@@ -40,7 +40,6 @@ import AWSLambdaRuntime
 public actor BreezeLambdaAPI<T: BreezeCodable>: Service {
 
     let logger = Logger(label: "service-group-breeze-lambda-api")
-    let timeout: TimeAmount
     private let serviceGroup: ServiceGroup
     private let apiConfig: any APIConfiguring
     private let dynamoDBService: BreezeDynamoDBService
@@ -57,10 +56,9 @@ public actor BreezeLambdaAPI<T: BreezeCodable>: Service {
     public init(apiConfig: APIConfiguring = BreezeAPIConfiguration()) async throws {
         do {
             self.apiConfig = apiConfig
-            self.timeout = .seconds(apiConfig.dbTimeout)
             let config = try apiConfig.getConfig()
             let httpConfig = BreezeHTTPClientConfig(
-                timeout: .seconds(60),
+                timeout: .seconds(apiConfig.dbTimeout),
                 logger: logger
             )
             let operation = try apiConfig.operation()

Tests/BreezeLambdaAPITests/BreezeLambdaAPITests.swift

@@ -18,9 +18,11 @@ import Testing
 import ServiceLifecycle
 import ServiceLifecycleTestKit
 import BreezeDynamoDBService
+import Foundation
+import AsyncHTTPClient
 
 struct APIConfiguration: APIConfiguring {
-    var dbTimeout: Int64 = 30
+    var dbTimeout: Int64 = 1
 
     func operation() throws -> BreezeOperation {
         .list
@@ -34,21 +36,16 @@ struct APIConfiguration: APIConfiguring {
 struct BreezeLambdaAPITests {
 
     let logger = Logger(label: "BreezeHTTPClientServiceTests")
-    
+
     @Test
     func test_breezeLambdaAPI_whenValidEnvironment() async throws {
-        do {
-        try await testGracefulShutdown { gracefulShutdownTestTrigger in
+        await testGracefulShutdown { gracefulShutdownTestTrigger in
             let (gracefulStream, continuation) = AsyncStream<Void>.makeStream()
-            try await withThrowingTaskGroup(of: Void.self) { group in
-                let sut = try await BreezeLambdaAPI<Product>(apiConfig: APIConfiguration())
+            await withThrowingTaskGroup(of: Void.self) { group in
                 group.addTask {
-                    try await withGracefulShutdownHandler{
-                        try await sut.run()
-                    } onGracefulShutdown: {
-                        logger.info("On Graceful Shutdown")
-                        continuation.yield()
-                    }
+                    let sut = try await BreezeLambdaAPI<Product>(apiConfig: APIConfiguration())
+                    try await sut.run()
+                    continuation.yield()
                 }
                 group.addTask {
                     try await Task.sleep(nanoseconds: 1_000_000_000)
@@ -61,36 +58,25 @@ struct BreezeLambdaAPITests {
                 }
             }
         }
-        } catch {
-            logger.error("Error during test: \(error)")
-            throw error
-        }
     }
 
     @Test
     func test_breezeLambdaAPI_whenInvalidEnvironment() async throws {
         await #expect(throws: BreezeLambdaAPIError.self) {
-            let (errorStream, continuation) = AsyncStream<Error>.makeStream()
+            let (errorStream, continuation) = AsyncStream<Void>.makeStream()
             try await testGracefulShutdown { gracefulShutdownTestTrigger in
+                let sut = try await BreezeLambdaAPI<Product>()
                 try await withThrowingTaskGroup(of: Void.self) { group in
-                    let sut = try await BreezeLambdaAPI<Product>()
                     group.addTask {
-                        try await withGracefulShutdownHandler {
-                            do {
-                                try await sut.run()
-                            } catch {
-                                continuation.yield(error)
-                                continuation.finish()
-                            }
-                        } onGracefulShutdown: {
-                            logger.info("Performing onGracefulShutdown")
-                        }
+                        try? await sut.run()
+                        continuation.yield()            
                     }
                     group.addTask {
                         try await Task.sleep(nanoseconds: 1_000_000_000)
                         gracefulShutdownTestTrigger.triggerGracefulShutdown()
                     }
                     for await _ in errorStream {
+                        continuation.finish()
                         logger.info("Error stream received")
                         group.cancelAll()
                     }