@@ -41,39 +41,78 @@ struct LambdaRuntimeTests {
41
41
)
42
42
43
43
try await withThrowingTaskGroup ( of: Void . self) { taskGroup in
44
+
44
45
// start the first runtime
45
46
taskGroup. addTask {
46
- // ChannelError will be thrown when we cancel the task group
47
- await #expect( throws: ChannelError . self) {
48
- try await runtime1. run ( )
49
- }
47
+ // will throw LambdaRuntimeError when run() is called second or ChannelError when cancelled
48
+ try await runtime1. run ( )
50
49
}
51
50
52
51
// wait a small amount to ensure runtime1 task is started
53
52
try await Task . sleep ( for: . seconds( 0.5 ) )
54
53
55
- // Running the second runtime should trigger LambdaRuntimeError
56
- await #expect( throws: LambdaRuntimeError . self) {
54
+ // start the second runtime
55
+ taskGroup. addTask {
56
+ // will throw LambdaRuntimeError when run() is called second or ChannelError when cancelled
57
57
try await runtime2. run ( )
58
58
}
59
59
60
- // cancel runtime 1 / task 1
60
+ // get the first result (should throw a LambdaRuntimeError)
61
+ try await #require( throws: LambdaRuntimeError . self) {
62
+ try await taskGroup. next ( )
63
+ }
64
+
65
+ // cancel the group to end the test
61
66
taskGroup. cancelAll ( )
67
+
62
68
}
69
+ }
70
+ @Test ( " run() must be cancellable " )
71
+ func testLambdaRuntimeCancellable( ) async throws {
63
72
64
- // Running the second runtime should work now
65
- try await withThrowingTaskGroup ( of: Void . self) { taskGroup in
66
- taskGroup. addTask {
67
- // ChannelError will be thrown when we cancel the task group
68
- await #expect( throws: ChannelError . self) {
69
- try await runtime2. run ( )
73
+ let logger = Logger ( label: " LambdaRuntimeTests.RuntimeCancellable " )
74
+ // create a runtime
75
+ let runtime = LambdaRuntime (
76
+ handler: MockHandler ( ) ,
77
+ eventLoop: Lambda . defaultEventLoop,
78
+ logger: logger
79
+ )
80
+
81
+ // Running the runtime with structured concurrency
82
+ // Task group returns when all tasks are completed.
83
+ // Even cancelled tasks must cooperatlivly complete
84
+ await #expect( throws: Never . self) {
85
+ try await withThrowingTaskGroup ( of: Void . self) { taskGroup in
86
+ taskGroup. addTask {
87
+ logger. trace ( " --- launching runtime ---- " )
88
+ try await runtime. run ( )
70
89
}
71
- }
72
90
73
- // Set timeout and cancel the runtime 2
74
- try await Task . sleep ( for: . seconds( 1 ) )
75
- taskGroup. cancelAll ( )
91
+ // Add a timeout task to the group
92
+ taskGroup. addTask {
93
+ logger. trace ( " --- launching timeout task ---- " )
94
+ try await Task . sleep ( for: . seconds( 5 ) )
95
+ if Task . isCancelled { return }
96
+ logger. trace ( " --- throwing timeout error ---- " )
97
+ throw TestError . timeout // Fail the test if the timeout triggers
98
+ }
99
+
100
+ do {
101
+ // Wait for the runtime to start
102
+ logger. trace ( " --- waiting for runtime to start ---- " )
103
+ try await Task . sleep ( for: . seconds( 1 ) )
104
+
105
+ // Cancel all tasks, this should not throw an error
106
+ // and should allow the runtime to complete gracefully
107
+ logger. trace ( " --- cancel all tasks ---- " )
108
+ taskGroup. cancelAll ( ) // Cancel all tasks
109
+ } catch {
110
+ logger. error ( " --- catch an error: \( error) " )
111
+ throw error // Propagate the error to fail the test
112
+ }
113
+ }
76
114
}
115
+
77
116
}
78
117
}
79
118
@@ -86,3 +125,15 @@ struct MockHandler: StreamingLambdaHandler {
86
125
87
126
}
88
127
}
128
+
129
+ // Define a custom error for timeout
130
+ enum TestError : Error , CustomStringConvertible {
131
+ case timeout
132
+
133
+ var description : String {
134
+ switch self {
135
+ case . timeout:
136
+ return " Test timed out waiting for the task to complete. "
137
+ }
138
+ }
139
+ }
0 commit comments