Additional Test Helpers

Package.swift

@@ -40,7 +40,8 @@ let targets: [Target] = [
     .target(
         name: "LambdaMocks",
         dependencies: [
-            "LambdaExtrasCore"
+            "LambdaExtrasCore",
+            .product(name: "NIO", package: "swift-nio")
         ]
     ),
     .testTarget(

Sources/LambdaExtras/Extensions.swift

@@ -11,16 +11,29 @@ import Foundation
 import LambdaExtrasCore
 import NIOCore
 
+extension Lambda {
+    /// Returns the value of the environment variable with the given name.
+    ///
+    /// This method throws ``EventHandler.envError`` if a value for the given environment variable
+    /// name is not found.
+    ///
+    /// - Parameter name: The name of the environment variable to return.
+    /// - Returns: The value of the given environment variable.
+    static func env(name: String) throws -> String {
+        guard let value = env(name) else {
+            throw HandlerError.envError(name)
+        }
+
+        return value
+    }
+}
+
 public extension EnvironmentValueProvider where EnvironmentVariable == String {
     /// Returns the value of the given environment variable.
     ///
     /// - Parameter environmentVariable: The environment variable whose value should be returned.
     func value(for environmentVariable: EnvironmentVariable) throws -> String {
-        guard let value = Lambda.env(environmentVariable) else {
-            throw HandlerError.envError(environmentVariable)
-        }
-
-        return value
+        try Lambda.env(name: environmentVariable)
     }
 }
 
@@ -29,11 +42,7 @@ public extension EnvironmentValueProvider where EnvironmentVariable: RawRepresen
     ///
     /// - Parameter environmentVariable: The environment variable whose value should be returned.
     func value(for environmentVariable: EnvironmentVariable) throws -> String {
-        guard let value = Lambda.env(environmentVariable.rawValue) else {
-            throw HandlerError.envError(environmentVariable.rawValue)
-        }
-
-        return value
+        try Lambda.env(name: environmentVariable.rawValue)
     }
 }
 

Sources/LambdaExtrasCore/Protocols/RuntimeContext.swift

@@ -40,4 +40,7 @@ public protocol RuntimeContext: Sendable {
 
     /// `ByteBufferAllocator` to allocate `ByteBuffer`.
     var allocator: ByteBufferAllocator { get }
+
+    /// Returns the time remaining before the deadline.
+    func getRemainingTime() -> TimeAmount
 }

Sources/LambdaMocks/ContextProvider.swift

@@ -0,0 +1,122 @@
+//
+//  ContextProvider.swift
+//  LambdaExtras
+//
+//  Created by Mathew Gacy on 1/7/24.
+//
+
+import Foundation
+import Logging
+import NIOCore
+import NIO
+
+/// A helper to create and manage mock initialization and runtime contexts for testing.
+///
+/// Example usage:
+///
+/// ```swift
+/// final class MyHandlerTests: XCTestCase {
+///     var contextProvider: ContextProvider<MyEnvironment>
+///
+///     override func setUp() {
+///         contextProvider.setUp()
+///     }
+///
+///     override func tearDown() {
+///         XCTAssertNoThrow(try contextProvider.shutdown())
+///     }
+///
+///     func testMyHandler() async throws {
+///         let sut = try await MyHandler(context: contextProvider.makeInitializationContext())
+///         let actual = try await sut.handle(MockEvent(), context: contextProvider.makeContext())
+///         ...
+///     }
+/// }
+/// ```
+public struct ContextProvider<EnvironmentVariable> {
+    /// The event loop group used to provide the contexts' event loops.
+    public private(set) var eventLoopGroup: EventLoopGroup!
+
+    /// The event loop for the contexts.
+    public private(set) var eventLoop: EventLoop!
+
+    /// The logger for the contexts.
+    public var logger: Logger
+
+    /// A closure returning the value of the given environment variable.
+    public var environmentValueProvider: @Sendable (EnvironmentVariable) throws -> String
+
+    /// Creates an instance.
+    ///
+    /// - Parameter environmentValueProvider: A closure returning the value of the given
+    /// environment variable.
+    public init(
+        logger: Logger = .mock,
+        environmentValueProvider: @escaping @Sendable (EnvironmentVariable) throws -> String
+    ) {
+        self.logger = logger
+        self.environmentValueProvider = environmentValueProvider
+    }
+
+    /// Sets up the event loop used for the provided initialization and runtime contexts.
+    ///
+    /// Call this in your test class's `setUp()` method:
+    ///
+    /// ```swift
+    /// final class MyHandlerTests: XCTestCase {
+    ///     var contextProvider: ContextProvider<MyEnvironment>
+    ///     ...
+    ///     override func setUp() {
+    ///         contextProvider.setUp()
+    ///         ...
+    ///     }
+    /// }
+    /// ```
+    public mutating func setUp() {
+        eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        eventLoop = eventLoopGroup.next()
+    }
+
+    /// Shuts the event loop group down.
+    ///
+    /// Call this in your test class's `.tearDown()` method:
+    ///
+    /// ```swift
+    /// final class MyHandlerTests: XCTestCase {
+    ///     var contextProvider: ContextProvider<MyEnvironment>
+    ///     ...
+    ///     override func tearDown() {
+    ///         XCTAssertNoThrow(try contextProvider.shutdown())
+    ///         ...
+    ///     }
+    /// }
+    /// ```
+    public mutating func shutdown() throws {
+        defer {
+            eventLoop = nil
+            eventLoopGroup = nil
+        }
+        try eventLoopGroup.syncShutdownGracefully()
+    }
+
+    /// Returns the mocked initialization context.
+    public func makeInitializationContext() -> MockInitializationContext<EnvironmentVariable> {
+        .init(
+            logger: logger,
+            eventLoop: eventLoop,
+            allocator: .init(),
+            environmentValueProvider: environmentValueProvider)
+    }
+
+    /// Returns the mocked runtime context.
+    ///
+    /// - Parameter configuration: The configuration for the mocked runtime context.
+    public func makeContext(
+        configuration: MockContext<EnvironmentVariable>.Configuration = .init()
+    ) -> MockContext<EnvironmentVariable> {
+        .init(
+            eventLoop: eventLoop,
+            configuration: configuration,
+            environmentValueProvider: environmentValueProvider)
+    }
+}

Sources/LambdaMocks/Extensions/Dispatch+Utils.swift

@@ -0,0 +1,15 @@
+//
+//  Dispatch+Utils.swift
+//  LambdaExtras
+//
+//  Created by Mathew Gacy on 1/19/24.
+//
+
+import Dispatch
+
+extension DispatchWallTime {
+    /// The interval between the point and its reference point.
+    var millisecondsSinceEpoch: Int64 {
+        Int64(bitPattern: self.rawValue) / -1_000_000
+    }
+}

Sources/LambdaMocks/Extensions/Logger+Utils.swift

@@ -0,0 +1,16 @@
+//
+//  Logger+Utils.swift
+//  LambdaExtras
+//
+//  Created by Mathew Gacy on 1/7/24.
+//
+
+import Foundation
+import Logging
+
+public extension Logger {
+    /// A logger for use in ``MockContext`` and ``MockInitializationContext``.
+    static let mock = Logger(
+        label: "mock-logger",
+        factory: { _ in StreamLogHandler.standardOutput(label: "mock-logger") })
+}

Sources/LambdaMocks/MockContext.swift

@@ -12,7 +12,7 @@ import Logging
 import NIOCore
 
 /// A mock function context for testing.
-public struct MockContext<E>: RuntimeContext, EnvironmentValueProvider {
+public struct MockContext<EnvironmentVariable>: RuntimeContext, EnvironmentValueProvider {
     public var requestID: String
     public var traceID: String
     public var invokedFunctionARN: String
@@ -23,8 +23,13 @@ public struct MockContext<E>: RuntimeContext, EnvironmentValueProvider {
     public var eventLoop: EventLoop
     public var allocator: ByteBufferAllocator
 
+    /// A closure returning a `TimeAmount` from a given `DispatchWallTime`.
+    ///
+    /// This is used to return the remaining time until the context's ``deadline``.
+    public var remainingTimeProvider: @Sendable (DispatchWallTime) -> TimeAmount
+
     /// A closure returning the value of the given environment variable.
-    public var environmentValueProvider: @Sendable (E) throws -> String
+    public var environmentValueProvider: @Sendable (EnvironmentVariable) throws -> String
 
     /// Creates a new instance.
     ///
@@ -38,6 +43,7 @@ public struct MockContext<E>: RuntimeContext, EnvironmentValueProvider {
     ///   - logger: The logger.
     ///   - eventLoop: The event loop.
     ///   - allocator: The byte buffer allocator.
+    ///   - remainingTimeProvider: A closure returning a `TimeAmount` from a given `DispatchWallTime`.
     ///   - environmentValueProvider: A closure returning the value of the given environment
     ///   variable.
     public init(
@@ -50,7 +56,8 @@ public struct MockContext<E>: RuntimeContext, EnvironmentValueProvider {
         logger: Logger,
         eventLoop: EventLoop,
         allocator: ByteBufferAllocator,
-        environmentValueProvider: @escaping @Sendable (E) throws -> String
+        remainingTimeProvider: @escaping @Sendable (DispatchWallTime) -> TimeAmount,
+        environmentValueProvider: @escaping @Sendable (EnvironmentVariable) throws -> String
     ) {
         self.requestID = requestID
         self.traceID = traceID
@@ -61,45 +68,104 @@ public struct MockContext<E>: RuntimeContext, EnvironmentValueProvider {
         self.logger = logger
         self.eventLoop = eventLoop
         self.allocator = allocator
+        self.remainingTimeProvider = remainingTimeProvider
         self.environmentValueProvider = environmentValueProvider
     }
 
-    public func value(for environmentVariable: E) throws -> String {
+    public func getRemainingTime() -> TimeAmount {
+        remainingTimeProvider(deadline)
+    }
+
+    public func value(for environmentVariable: EnvironmentVariable) throws -> String {
         try environmentValueProvider(environmentVariable)
     }
 }
 
 public extension MockContext {
+
+    /// Configuration data for ``MockContext``.
+    struct Configuration {
+        /// The request ID, which identifies the request that triggered the function invocation.
+        public var requestID: String
+
+        /// The AWS X-Ray tracing header.
+        public var traceID: String
+
+        /// The ARN of the Lambda function, version, or alias that's specified in the invocation.
+        public var invokedFunctionARN: String
+
+        /// The time interval before the context's deadline.
+        public var timeout: DispatchTimeInterval
+
+        /// For invocations from the AWS Mobile SDK, data about the Amazon Cognito identity provider.
+        public var cognitoIdentity: String?
+
+        /// For invocations from the AWS Mobile SDK, data about the client application and device.
+        public var clientContext: String?
+
+        /// Creates an instance.
+        ///
+        /// - Parameters:
+        ///   - requestID: The request ID.
+        ///   - traceID: The AWS X-Ray tracing header.
+        ///   - invokedFunctionARN: The ARN of the Lambda function.
+        ///   - timeout: The time interval before the context's deadline.
+        ///   - cognitoIdentity: Data about the Amazon Cognito identity provider.
+        ///   - clientContext: Data about the client application and device.
+        public init(
+            requestID: String = "\(DispatchTime.now().uptimeNanoseconds)",
+            traceID: String = "Root=\(DispatchTime.now().uptimeNanoseconds);Parent=\(DispatchTime.now().uptimeNanoseconds);Sampled=1",
+            invokedFunctionARN: String = "arn:aws:lambda:us-east-1:\(DispatchTime.now().uptimeNanoseconds):function:custom-runtime",
+            timeout: DispatchTimeInterval = .seconds(5),
+            cognitoIdentity: String? = nil,
+            clientContext: String? = nil
+        ) {
+            self.requestID = requestID
+            self.traceID = traceID
+            self.invokedFunctionARN = invokedFunctionARN
+            self.timeout = timeout
+            self.cognitoIdentity = cognitoIdentity
+            self.clientContext = clientContext
+        }
+    }
+
+    /// Returns the time interval between a given point in time and the current time.
+    ///
+    /// - Parameter deadline: The time with which to compare now.
+    /// - Returns: The time interval between the given deadline and now.
+    @Sendable
+    static func timeAmountUntil(_ deadline: DispatchWallTime) -> TimeAmount {
+        .milliseconds(deadline.millisecondsSinceEpoch - DispatchWallTime.now().millisecondsSinceEpoch)
+    }
+
     /// Creates a new instance.
     ///
     /// - Parameters:
-    ///   - timeout: The time interval at which the function will time out.
-    ///   - requestID: The request ID.
-    ///   - traceID: The tracing header.
-    ///   - invokedFunctionARN: The ARN of the Lambda function.
     ///   - eventLoop: The event loop.
+    ///   - configuration: The context configuration.
+    ///   - logger: The logger.
     ///   - allocator: The byte buffer allocator.
+    ///   - remainingTimeProvider:
     ///   - environmentValueProvider: A closure returning the value of the given environment
     ///   variable.
     init(
-        timeout: DispatchTimeInterval = .seconds(3),
-        requestID: String = UUID().uuidString,
-        traceID: String = "abc123",
-        invokedFunctionARN: String = "aws:arn:",
         eventLoop: EventLoop,
+        configuration: Configuration = .init(),
+        logger: Logger = .mock,
         allocator: ByteBufferAllocator = .init(),
-        environmentValueProvider: @escaping @Sendable (E) throws -> String
+        remainingTimeProvider: @escaping @Sendable (DispatchWallTime) -> TimeAmount = Self.timeAmountUntil,
+        environmentValueProvider: @escaping @Sendable (EnvironmentVariable) throws -> String
     ) {
-        self.requestID = requestID
-        self.traceID = traceID
-        self.invokedFunctionARN = invokedFunctionARN
-        self.deadline = .now() + timeout
-        self.logger = Logger(
-            label: "mock-logger",
-            factory: { _ in StreamLogHandler.standardOutput(label: "mock-logger") }
-        )
+        self.requestID = configuration.requestID
+        self.traceID = configuration.traceID
+        self.invokedFunctionARN = configuration.invokedFunctionARN
+        self.deadline = .now() + configuration.timeout
+        self.cognitoIdentity = configuration.cognitoIdentity
+        self.clientContext = configuration.clientContext
+        self.logger = logger
         self.eventLoop = eventLoop
         self.allocator = allocator
+        self.remainingTimeProvider = remainingTimeProvider
         self.environmentValueProvider = environmentValueProvider
     }
 }