diff --git a/bitchatTests/Fragmentation/FragmentationTests.swift b/bitchatTests/Fragmentation/FragmentationTests.swift index 1632f9ad..e507fe25 100644 --- a/bitchatTests/Fragmentation/FragmentationTests.swift +++ b/bitchatTests/Fragmentation/FragmentationTests.swift @@ -32,29 +32,28 @@ struct FragmentationTests { ) let capture = CaptureDelegate() ble.delegate = capture - + // Construct a big packet (3KB) from a remote sender (not our own ID) let remoteShortID = PeerID(str: "1122334455667788") let original = makeLargePublicPacket(senderShortHex: remoteShortID, size: 3_000) - + // Use a small fragment size to ensure multiple pieces let fragments = fragmentPacket(original, fragmentSize: 400) - + // Shuffle fragments to simulate out-of-order arrival let shuffled = fragments.shuffled() - - // Inject fragments spaced out to avoid concurrent mutation inside BLEService + + // Send fragments sequentially with small delays (no fire-and-forget Tasks) for (i, fragment) in shuffled.enumerated() { - let delay = 5 * Double(i) * 0.001 - Task { - try await sleep(delay) - ble._test_handlePacket(fragment, fromPeerID: remoteShortID) + if i > 0 { + try await Task.sleep(for: .milliseconds(5)) } + ble._test_handlePacket(fragment, fromPeerID: remoteShortID) } - - // Allow async processing - try await sleep(0.5) - + + // Wait for delegate callback with proper timeout + try await capture.waitForPublicMessages(count: 1, timeout: .seconds(2)) + #expect(capture.publicMessages.count == 1) #expect(capture.publicMessages.first?.content.count == 3_000) } @@ -68,26 +67,26 @@ struct FragmentationTests { ) let capture = CaptureDelegate() ble.delegate = capture - + let remoteShortID = PeerID(str: "A1B2C3D4E5F60708") let original = makeLargePublicPacket(senderShortHex: remoteShortID, size: 2048) var frags = fragmentPacket(original, fragmentSize: 300) - + // Duplicate one fragment if let dup = frags.first { frags.insert(dup, at: 1) } - + + // Send fragments sequentially with small delays (no fire-and-forget Tasks) for (i, fragment) in frags.enumerated() { - let delay = 5 * Double(i) * 0.001 - Task { - try await sleep(delay) - ble._test_handlePacket(fragment, fromPeerID: remoteShortID) + if i > 0 { + try await Task.sleep(for: .milliseconds(5)) } + ble._test_handlePacket(fragment, fromPeerID: remoteShortID) } - - // Allow async processing - try await sleep(0.5) + + // Wait for delegate callback with proper timeout + try await capture.waitForPublicMessages(count: 1, timeout: .seconds(2)) #expect(capture.publicMessages.count == 1) #expect(capture.publicMessages.first?.content.count == 2048) @@ -196,12 +195,128 @@ struct FragmentationTests { } extension FragmentationTests { - private final class CaptureDelegate: BitchatDelegate { - var publicMessages: [(peerID: PeerID, nickname: String, content: String)] = [] - var receivedMessages: [BitchatMessage] = [] - func didReceiveMessage(_ message: BitchatMessage) { - receivedMessages.append(message) + /// Thread-safe delegate that supports awaiting message delivery + private final class CaptureDelegate: BitchatDelegate, @unchecked Sendable { + private let lock = NSLock() + private var _publicMessages: [(peerID: PeerID, nickname: String, content: String)] = [] + private var _receivedMessages: [BitchatMessage] = [] + private var publicMessageContinuation: CheckedContinuation? + private var receivedMessageContinuation: CheckedContinuation? + private var expectedPublicMessageCount: Int = 0 + private var expectedReceivedMessageCount: Int = 0 + + var publicMessages: [(peerID: PeerID, nickname: String, content: String)] { + lock.lock() + defer { lock.unlock() } + return _publicMessages } + + var receivedMessages: [BitchatMessage] { + lock.lock() + defer { lock.unlock() } + return _receivedMessages + } + + func didReceiveMessage(_ message: BitchatMessage) { + lock.lock() + _receivedMessages.append(message) + let count = _receivedMessages.count + let expected = expectedReceivedMessageCount + let continuation = receivedMessageContinuation + lock.unlock() + + if count >= expected, let cont = continuation { + lock.lock() + receivedMessageContinuation = nil + lock.unlock() + cont.resume() + } + } + + func didReceivePublicMessage(from peerID: PeerID, nickname: String, content: String, timestamp: Date, messageID: String?) { + lock.lock() + _publicMessages.append((peerID, nickname, content)) + let count = _publicMessages.count + let expected = expectedPublicMessageCount + let continuation = publicMessageContinuation + lock.unlock() + + if count >= expected, let cont = continuation { + lock.lock() + publicMessageContinuation = nil + lock.unlock() + cont.resume() + } + } + + /// Waits for the specified number of public messages to be received + func waitForPublicMessages(count: Int, timeout: Duration = .seconds(2)) async throws { + lock.lock() + if _publicMessages.count >= count { + lock.unlock() + return + } + expectedPublicMessageCount = count + lock.unlock() + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + await withCheckedContinuation { continuation in + self.lock.lock() + // Recheck count after acquiring lock to avoid race condition + // where message arrives between initial check and continuation install + if self._publicMessages.count >= count { + self.lock.unlock() + continuation.resume() + return + } + self.publicMessageContinuation = continuation + self.lock.unlock() + } + } + group.addTask { + try await Task.sleep(for: timeout) + throw CancellationError() + } + try await group.next() + group.cancelAll() + } + } + + /// Waits for the specified number of received messages + func waitForReceivedMessages(count: Int, timeout: Duration = .seconds(2)) async throws { + lock.lock() + if _receivedMessages.count >= count { + lock.unlock() + return + } + expectedReceivedMessageCount = count + lock.unlock() + + try await withThrowingTaskGroup(of: Void.self) { group in + group.addTask { + await withCheckedContinuation { continuation in + self.lock.lock() + // Recheck count after acquiring lock to avoid race condition + // where message arrives between initial check and continuation install + if self._receivedMessages.count >= count { + self.lock.unlock() + continuation.resume() + return + } + self.receivedMessageContinuation = continuation + self.lock.unlock() + } + } + group.addTask { + try await Task.sleep(for: timeout) + throw CancellationError() + } + try await group.next() + group.cancelAll() + } + } + func didConnectToPeer(_ peerID: PeerID) {} func didDisconnectFromPeer(_ peerID: PeerID) {} func didUpdatePeerList(_ peers: [PeerID]) {} @@ -209,9 +324,6 @@ extension FragmentationTests { func didUpdateMessageDeliveryStatus(_ messageID: String, status: DeliveryStatus) {} func didReceiveNoisePayload(from peerID: PeerID, type: NoisePayloadType, payload: Data, timestamp: Date) {} func didUpdateBluetoothState(_ state: CBManagerState) {} - func didReceivePublicMessage(from peerID: PeerID, nickname: String, content: String, timestamp: Date, messageID: String?) { - publicMessages.append((peerID, nickname, content)) - } func didReceiveRegionalPublicMessage(from peerID: PeerID, nickname: String, content: String, timestamp: Date) {} }