mobile-ios/RosettaTests/AttachmentParityTests.swift

import UIKit
import XCTest
@testable import Rosetta

@MainActor
final class AttachmentParityTests: XCTestCase {
    private var ctx: DBTestContext!
    private var transportMock: MockAttachmentFlowTransport!
    private var senderMock: MockPacketFlowSender!

    private var ownPrivateKeyHex: String = ""
    private var ownPublicKey: String = ""
    private var peerPublicKey: String = ""

    override func setUpWithError() throws {
        let ownPair = try Self.makeKeyPair()
        let peerPair = try Self.makeKeyPair()

        ownPrivateKeyHex = ownPair.privateKeyHex
        ownPublicKey = ownPair.publicKeyHex
        peerPublicKey = peerPair.publicKeyHex

        ctx = DBTestContext(account: ownPublicKey)
        transportMock = MockAttachmentFlowTransport()
        senderMock = MockPacketFlowSender()

        SessionManager.shared.testConfigureSessionForParityFlows(
            currentPublicKey: ownPublicKey,
            privateKeyHex: ownPrivateKeyHex
        )
        SessionManager.shared.attachmentFlowTransport = transportMock
        SessionManager.shared.packetFlowSender = senderMock
        AttachmentCache.shared.privateKey = ownPrivateKeyHex
    }

    override func tearDownWithError() throws {
        ctx?.teardown()
        ctx = nil
        transportMock = nil
        senderMock = nil
        AttachmentCache.shared.privateKey = nil
        SessionManager.shared.testResetParityFlowDependencies()
    }

    func testAttachmentPreviewParserMatrix() {
        let tag = "aaaaaaaa-1111-2222-3333-bbbbbbbbbbbb"

        XCTAssertEqual(AttachmentPreviewCodec.downloadTag(from: "\(tag)::LKO2"), tag)
        XCTAssertEqual(AttachmentPreviewCodec.blurHash(from: "\(tag)::LKO2"), "LKO2")

        XCTAssertEqual(AttachmentPreviewCodec.downloadTag(from: "::LKO2"), "")
        XCTAssertEqual(AttachmentPreviewCodec.blurHash(from: "::LKO2"), "LKO2")

        let taggedFile = AttachmentPreviewCodec.parseFilePreview("\(tag)::2048::report.pdf")
        XCTAssertEqual(taggedFile.downloadTag, tag)
        XCTAssertEqual(taggedFile.fileSize, 2048)
        XCTAssertEqual(taggedFile.fileName, "report.pdf")

        let localFile = AttachmentPreviewCodec.parseFilePreview("512::notes.txt")
        XCTAssertEqual(localFile.downloadTag, "")
        XCTAssertEqual(localFile.fileSize, 512)
        XCTAssertEqual(localFile.fileName, "notes.txt")

        XCTAssertEqual(AttachmentPreviewCodec.payload(from: "legacy_preview"), "legacy_preview")
    }

    func testOutgoingAttachmentPacketShapeClearsBlobAndUsesTaggedPreview() async throws {
        try await ctx.bootstrap()

        let image = Self.makeSolidImage(color: .systemBlue)
        let imageAttachment = PendingAttachment.fromImage(image)
        let fileData = Data("hello parity".utf8)
        let fileAttachment = PendingAttachment.fromFile(data: fileData, fileName: "notes.txt")

        let imageTag = "11111111-1111-1111-1111-111111111111"
        let fileTag = "22222222-2222-2222-2222-222222222222"
        transportMock.tagsById[imageAttachment.id] = imageTag
        transportMock.tagsById[fileAttachment.id] = fileTag

        try await SessionManager.shared.sendMessageWithAttachments(
            text: "",
            attachments: [imageAttachment, fileAttachment],
            toPublicKey: peerPublicKey,
            opponentTitle: "Peer",
            opponentUsername: "peer"
        )

        XCTAssertEqual(transportMock.uploadedIds.count, 2)
        XCTAssertEqual(senderMock.sentMessages.count, 1)

        guard let sent = senderMock.sentMessages.first else {
            XCTFail("No outgoing packet captured")
            return
        }
        XCTAssertEqual(sent.attachments.count, 2)
        XCTAssertTrue(sent.attachments.allSatisfy { $0.blob.isEmpty })

        guard let sentImage = sent.attachments.first(where: { $0.id == imageAttachment.id }) else {
            XCTFail("Missing image attachment in packet")
            return
        }
        XCTAssertEqual(sentImage.transportTag, imageTag)
        XCTAssertEqual(AttachmentPreviewCodec.downloadTag(from: sentImage.preview), "")

        guard let sentFile = sent.attachments.first(where: { $0.id == fileAttachment.id }) else {
            XCTFail("Missing file attachment in packet")
            return
        }
        XCTAssertEqual(sentFile.transportTag, fileTag)
        let parsedFile = AttachmentPreviewCodec.parseFilePreview(sentFile.preview)
        XCTAssertEqual(parsedFile.downloadTag, "")
        XCTAssertEqual(parsedFile.fileSize, fileData.count)
        XCTAssertEqual(parsedFile.fileName, "notes.txt")
    }

    func testSavedSelfFileFlowKeepsLocalFileOpenableWithoutUpload() async throws {
        try await ctx.bootstrap()

        let fileData = Data("self file payload".utf8)
        let fileAttachment = PendingAttachment.fromFile(data: fileData, fileName: "local.txt")

        try await SessionManager.shared.sendMessageWithAttachments(
            text: "",
            attachments: [fileAttachment],
            toPublicKey: ownPublicKey
        )

        XCTAssertTrue(transportMock.uploadedIds.isEmpty)
        XCTAssertTrue(senderMock.sentMessages.isEmpty)

        let cachedURL = AttachmentCache.shared.fileURL(
            forAttachmentId: fileAttachment.id,
            fileName: "local.txt"
        )
        XCTAssertNotNil(cachedURL)

        let loaded = AttachmentCache.shared.loadFileData(
            forAttachmentId: fileAttachment.id,
            fileName: "local.txt"
        )
        XCTAssertEqual(loaded, fileData)
    }
}

private extension AttachmentParityTests {
    static func makeSolidImage(color: UIColor) -> UIImage {
        let renderer = UIGraphicsImageRenderer(size: CGSize(width: 32, height: 32))
        return renderer.image { ctx in
            color.setFill()
            ctx.fill(CGRect(x: 0, y: 0, width: 32, height: 32))
        }
    }

    static func makeKeyPair() throws -> (privateKeyHex: String, publicKeyHex: String) {
        let mnemonic = try CryptoManager.shared.generateMnemonic()
        let pair = try CryptoManager.shared.deriveKeyPair(from: mnemonic)
        return (pair.privateKey.hexString, pair.publicKey.hexString)
    }
}

private final class MockAttachmentFlowTransport: AttachmentFlowTransporting {
    var tagsById: [String: String] = [:]
    private(set) var uploadedIds: [String] = []

    func uploadFile(id: String, content: Data) async throws -> (tag: String, server: String) {
        uploadedIds.append(id)
        let tag = tagsById[id] ?? UUID().uuidString.lowercased()
        return (tag: tag, server: "https://mock-transport.test")
    }

    func downloadFile(tag: String, server: String?) async throws -> Data {
        Data()
    }
}

private final class MockPacketFlowSender: PacketFlowSending {
    private(set) var sentMessages: [PacketMessage] = []

    func sendPacket(_ packet: any Packet) {
        if let message = packet as? PacketMessage {
            sentMessages.append(message)
        }
    }
}