SignalAttachmentTest.swift

//
// Copyright 2021 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//

import CoreServices
import SignalServiceKit
import UniformTypeIdentifiers
import XCTest
@testable import SignalUI

class SignalAttachmentTest: XCTestCase {

    // MARK: - Utilities

    func testMetadataStrippingDoesNotChangeOrientation(url: URL) throws {
        let size = (try? DataImageSource.forPath(url.path).imageMetadata()?.pixelSize) ?? .zero

        let dataSource = DataSourcePath(fileUrl: url, ownership: .borrowed)
        let attachment = try PreviewableAttachment.imageAttachment(
            dataSource: dataSource,
            dataUTI: UTType.jpeg.identifier,
        )
        let newSize = DataImageSource(try attachment.dataSource.readData()).imageMetadata()?.pixelSize

        XCTAssertEqual(newSize, size, "image dimensions changed for \(url.lastPathComponent)")
    }

    private func pngChunks(data: Data) throws -> [PngChunker.Chunk] {
        let chunker = try PngChunker(source: DataImageSource(data))
        var result = [PngChunker.Chunk]()
        while let chunk = try chunker.next() {
            result.append(chunk)
        }
        return result
    }

    private func pngChunkTypes(data: Data) throws -> [String] {
        (try pngChunks(data: data)).compactMap { chunk in
            String(data: chunk.type, encoding: .ascii)
        }
    }

    // MARK: - Tests

    func testMetadataStrippingDoesNotChangeOrientation() throws {
        let testBundle = Bundle(for: Self.self)
        try testMetadataStrippingDoesNotChangeOrientation(url: testBundle.url(
            forResource: "test-jpg",
            withExtension: "jpg",
        )!)
        try testMetadataStrippingDoesNotChangeOrientation(url: testBundle.url(
            forResource: "test-jpg-rotated",
            withExtension: "jpg",
        )!)
    }

    func testRemoveMetadataFromPng() throws {
        let testBundle = Bundle(for: Self.self)
        let url = testBundle.url(forResource: "test-png-with-metadata", withExtension: "png")!
        let dataSource = DataSourcePath(fileUrl: url, ownership: .borrowed)
        XCTAssertEqual(
            try pngChunkTypes(data: dataSource.readData()),
            ["IHDR", "PLTE", "sRGB", "tIME", "tEXt", "IDAT", "IEND"],
            "Test is not set up correctly. Fixture doesn't have the expected chunks",
        )

        let normalizedImage = try NormalizedImage.forDataSource(dataSource, dataUTI: UTType.png.identifier)
        let finalizedImage = try normalizedImage.finalizeImage(imageQuality: .three)

        XCTAssertEqual(
            try pngChunkTypes(data: try finalizedImage.dataSource.readData()),
            ["IHDR", "PLTE", "sRGB", "IDAT", "IEND"],
        )
    }

    func testRemoveMetadataFromAPng() throws {
        let pngData: Data = {
            let testBundle = Bundle(for: Self.self)
            let apngUrl = testBundle.url(forResource: "test-apng", withExtension: "png")!
            let apngData = try! Data(contentsOf: apngUrl)
            let apngChunks = try! pngChunks(data: apngData)

            // This is a `tEXt` chunk with some sample data.
            let newChunkData = Data([
                0x00,
                0x00,
                0x00,
                0x14,
                0x74,
                0x45,
                0x58,
                0x74,
                0x43,
                0x6f,
                0x6d,
                0x6d,
                0x65,
                0x6e,
                0x74,
                0x00,
                0x48,
                0x65,
                0x6c,
                0x6c,
                0x6f,
                0x20,
                0x77,
                0x6f,
                0x72,
                0x6c,
                0x64,
                0x21,
                0xab,
                0x32,
                0xb0,
                0x28,
            ])
            let newChunks: [Data] = (
                apngChunks.dropLast(1).map { $0.allBytes() } +
                    [newChunkData, apngChunks.last!.allBytes()],
            )

            return PngChunker.pngSignature + newChunks.reduce(Data()) { $0 + $1 }
        }()
        XCTAssert(
            (try pngChunkTypes(data: pngData)).contains("tEXt"),
            "Test is not set up correctly. Fixture doesn't have the expected chunks",
        )
        let dataSource = try DataSourcePath(writingTempFileData: pngData, fileExtension: "png")

        let normalizedImage = try NormalizedImage.forDataSource(dataSource, dataUTI: UTType.png.identifier)
        let finalizedImage = try normalizedImage.finalizeImage(imageQuality: .three)

        XCTAssert(
            !(try pngChunkTypes(data: try finalizedImage.dataSource.readData())).contains("tEXt"),
            "Result contained unexpected chunk",
        )
    }
}