macsync/@packages/imessage/Tests/IMessageSyncTests/ChunkingTests.swift

import Foundation
import Testing
@testable import IMessageSync

private func makeTestMessage(
    guid: String = "test-guid",
    text: String? = "Hello",
    attachments: [iMessageAttachment] = [],
    expressiveSendStyleId: String? = nil,
    associatedMessageGuid: String? = nil,
    senderDisplayName: String? = nil
) -> iMessage {
    iMessage(
        guid: guid,
        conversationId: "conv-001",
        senderId: nil,
        isFromMe: true,
        date: Date(),
        dateDelivered: nil,
        dateRead: nil,
        text: text,
        attributedBody: nil,
        associatedMessageType: nil,
        associatedMessageGuid: associatedMessageGuid,
        isAudioMessage: false,
        expressiveSendStyleId: expressiveSendStyleId,
        replyToGuid: nil,
        threadOriginatorGuid: nil,
        groupTitle: nil,
        balloonBundleId: nil,
        service: "iMessage",
        senderIdentifier: nil,
        senderDisplayName: senderDisplayName,
        senderPhoneNumber: nil,
        senderEmail: nil,
        attachments: attachments,
        attachmentsCount: attachments.count,
        attachmentsTotalSize: attachments.reduce(0) { $0 + $1.size },
        attachmentsFiletypes: Array(Set(attachments.compactMap { $0.mimeType }))
    )
}

@Suite("Adaptive Chunking")
struct AdaptiveChunkingTests {

    @Test func emptyMessagesProducesEmptyChunks() {
        let chunks = ChunkingStrategy.createAdaptiveChunks([])
        #expect(chunks.isEmpty)
    }

    @Test func singleSmallMessageProducesOneChunk() {
        let chunks = ChunkingStrategy.createAdaptiveChunks([makeTestMessage()])
        #expect(chunks.count == 1)
        #expect(chunks[0].count == 1)
    }

    @Test func multipleSmallMessagesUnderLimitProduceSingleChunk() {
        let messages = (0..<10).map { makeTestMessage(guid: "msg-\($0)") }
        let chunks = ChunkingStrategy.createAdaptiveChunks(messages)
        #expect(chunks.count == 1)
        #expect(chunks[0].count == 10)
    }

    @Test func messagesExceedingLimitSplitIntoMultipleChunks() {
        let largeData = String(repeating: "A", count: 60_000_000)
        let attachment = iMessageAttachment(
            filename: "large.bin",
            mimeType: "application/octet-stream",
            transferName: "large.bin",
            size: 60_000_000,
            data: largeData
        )
        let messages = (0..<3).map {
            makeTestMessage(guid: "large-\($0)", attachments: [attachment])
        }
        let chunks = ChunkingStrategy.createAdaptiveChunks(messages)
        #expect(chunks.count >= 2)
    }

    @Test func singleOversizedMessageGetsOwnChunk() {
        let hugeData = String(repeating: "X", count: 110_000_000)
        let attachment = iMessageAttachment(
            filename: "huge.bin",
            mimeType: "application/octet-stream",
            transferName: "huge.bin",
            size: 110_000_000,
            data: hugeData
        )
        let messages = [
            makeTestMessage(guid: "small-1"),
            makeTestMessage(guid: "huge", attachments: [attachment]),
            makeTestMessage(guid: "small-2"),
        ]
        let chunks = ChunkingStrategy.createAdaptiveChunks(messages)
        #expect(chunks.count == 3)
        #expect(chunks[0].count == 1)
        #expect(chunks[0][0].guid == "small-1")
        #expect(chunks[1].count == 1)
        #expect(chunks[1][0].guid == "huge")
        #expect(chunks[2].count == 1)
        #expect(chunks[2][0].guid == "small-2")
    }

    @Test func baseOverheadIsAccounted() {
        #expect(ChunkingStrategy.baseOverhead == 200)
    }

    @Test func maxChunkSizeBytesIs100MB() {
        #expect(ChunkingStrategy.maxChunkSizeBytes == 100_000_000)
    }
}

@Suite("Message Size Estimation")
struct MessageSizeEstimationTests {

    @Test func baseOverheadIs500Bytes() {
        let size = ChunkingStrategy.estimateMessageSize(makeTestMessage(text: nil))
        #expect(size == 500)
    }

    @Test func textFieldAddsUtf8ByteCount() {
        let text = "Hello, World!"
        let size = ChunkingStrategy.estimateMessageSize(makeTestMessage(text: text))
        #expect(size == 500 + text.utf8.count)
    }

    @Test func attachmentDataAddsUtf8ByteCount() {
        let data = "base64encodedcontent"
        let attachment = iMessageAttachment(
            filename: "test.jpg",
            mimeType: "image/jpeg",
            transferName: "test.jpg",
            size: 1000,
            data: data
        )
        let size = ChunkingStrategy.estimateMessageSize(makeTestMessage(text: nil, attachments: [attachment]))
        let expected = 500 + 100 + data.utf8.count + "test.jpg".utf8.count + "image/jpeg".utf8.count + "test.jpg".utf8.count
        #expect(size == expected)
    }

    @Test func multipleAttachmentsSumCorrectly() {
        let att1 = iMessageAttachment(filename: "a.jpg", mimeType: "image/jpeg", transferName: nil, size: 100, data: "AAA")
        let att2 = iMessageAttachment(filename: "b.png", mimeType: "image/png", transferName: nil, size: 200, data: "BBB")
        let size = ChunkingStrategy.estimateMessageSize(makeTestMessage(text: nil, attachments: [att1, att2]))
        let expected = 500
            + 100 + "AAA".utf8.count + "a.jpg".utf8.count + "image/jpeg".utf8.count
            + 100 + "BBB".utf8.count + "b.png".utf8.count + "image/png".utf8.count
        #expect(size == expected)
    }

    @Test func nilOptionalFieldsAddZero() {
        let size = ChunkingStrategy.estimateMessageSize(
            makeTestMessage(text: nil, expressiveSendStyleId: nil, associatedMessageGuid: nil, senderDisplayName: nil)
        )
        #expect(size == 500)
    }

    @Test func optionalStringFieldsAddToSize() {
        let size = ChunkingStrategy.estimateMessageSize(
            makeTestMessage(
                text: nil,
                expressiveSendStyleId: "com.apple.MobileSMS.expressivesend.impact",
                senderDisplayName: "Alice"
            )
        )
        let expected = 500
            + "com.apple.MobileSMS.expressivesend.impact".utf8.count
            + "Alice".utf8.count
        #expect(size == expected)
    }
}