Fix bug in the Swift Ice/ami test

swift/test/Ice/ami/AllTests.swift

@@ -4,9 +4,6 @@ import Foundation
 import Ice
 import TestCommon
 
-// WORKAROUND: Disable optimization for this test. Otherwise we get a crash in
-// "testing batch requests with communicator"
-@_optimize(none)
 func allTests(_ helper: TestHelper, collocated: Bool = false) async throws {
     func test(_ value: Bool, file: String = #file, line: Int = #line) throws {
         try helper.test(value, file: file, line: line)

swift/test/Ice/ami/TestI.swift

@@ -125,22 +125,26 @@ class TestI: TestIntf {
     }
 
     func waitForBatch(count: Int32, current _: Current) async throws -> Bool {
-        if await count == _state.getBatchCount() {
-            await _state.clearBatchCount()
-            return true
-        }
-
         var sub: AnyCancellable?
-        let currentCount = await withCheckedContinuation { continuation in
+        await withCheckedContinuation { continuation in
             sub = _batchSubject.sink { batchCount in
-                if batchCount >= count {
-                    continuation.resume(returning: batchCount)
+                if batchCount == count {
+                    continuation.resume()
+                } else if batchCount > count {
+                    fatalError("Received more batches than expected: \(batchCount) > \(count)")
                 }
             }
         }
         sub!.cancel()
-        let result = count == currentCount
+        // Check the batch count again after the sink has been cancelled to get the current value.
+        // The client only ever sends the expect number of batches so it is impossible that the
+        // batch count will be greater than the expected count at this point unless there is a bug.
+        // The above sink with also crash if the continuation is resumed twice if an unexpected batch is
+        // received before the subscription is cancelled.
+        let batchCount = await _state.getBatchCount()
+        let result = count == batchCount
         await _state.clearBatchCount()
+        _batchSubject.send(0) // Reset batch subject
         return result
     }