Repository: apple/swift-nio
File: Tests/NIOCoreTests/AsyncSequences/NIOThrowingAsyncSequenceTests.swift (Swift)

Code before:

    func testIteratorThrows_whenCancelled() async {
        let maxSequenceValue = 99
        _ = self.source.yield(contentsOf: Array(0...maxSequenceValue))
        await withThrowingTaskGroup(of: Void.self) { group in
            group.addTask {
                var counter = 0
                guard let sequence = self.sequence else {
                    return XCTFail("Expected to have an AsyncSequence")
                }

                do {
                    for try await next in sequence {
                        XCTAssertEqual(next, counter)
                        
                        if next < maxSequenceValue {
                            // This loop will exit when the task is cancelled
                            // or when all of the elements have been iterated.
                            // It is possible that the cancellation will be
                            // triggered after the last element was yielded,
                            // just before the iterator can return the `nil`
                            // signifying the end of the sequence.
                            // If this happens, we avoid incrementing the counter,
                            // so that we don't increase it to a value beyond the
                            // maximum possible sequence value.
                            counter += 1
                        }
                    }
                    XCTFail("Expected that this throws")
                } catch is CancellationError {
                    // expected
                } catch {
                    XCTFail("Unexpected error: \(error)")
                }

                XCTAssertLessThan(counter, 100)
            }

            group.cancelAll()
        }

        XCTAssertEqualWithoutAutoclosure(await self.delegate.events.prefix(1).collect(), [.didTerminate])
    }
}

// This is needed until async let is supported to be used in autoclosures
fileprivate func XCTAssertEqualWithoutAutoclosure<T>(
    _ expression1: T,
    _ expression2: T,
    _ message: @autoclosure () -> String = "",
    file: StaticString = #filePath,

Review comment (suggestion):
I don't quite follow that change nor the comment. The counter increment is happening inside the `for await` loop. The `AsyncSequence` that we are iterating over has exactly 100 elements and regardless when we are cancelled we should never count to above the element size of the `Array`.

Code after:

    func testIteratorThrows_whenCancelled() async {
        _ = self.source.yield(contentsOf: Array(1...100))
        await withThrowingTaskGroup(of: Void.self) { group in
            group.addTask {
                var itemsYieldedCounter = 0
                guard let sequence = self.sequence else {
                    return XCTFail("Expected to have an AsyncSequence")
                }

                do {
                    for try await next in sequence {
                        itemsYieldedCounter += 1
                        XCTAssertEqual(next, itemsYieldedCounter)
                    }
                    XCTFail("Expected that this throws")
                } catch is CancellationError {
                    // expected
                } catch {
                    XCTFail("Unexpected error: \(error)")
                }

                XCTAssertLessThanOrEqual(itemsYieldedCounter, 100)
            }

            group.cancelAll()
        }

        XCTAssertEqualWithoutAutoclosure(await self.delegate.events.prefix(1).collect(), [.didTerminate])
    }
}

// This is needed until async let is supported to be used in autoclosures
fileprivate func XCTAssertEqualWithoutAutoclosure<T>(
    _ expression1: T,
    _ expression2: T,
    _ message: @autoclosure () -> String = "",
    file: StaticString = #filePath,
    line: UInt = #line
) where T: Equatable {
    XCTAssertEqual(expression1, expression2, message(), file: file, line: line)
}

@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
extension AsyncSequence {
    /// Collect all elements in the sequence into an array.
    fileprivate func collect() async rethrows -> [Element] {
        try await self.reduce(into: []) { accumulated, next in
            accumulated.append(next)
        }
    }