AudioSnapshotTestingTests.swift

import Testing
import AVFAudio
@testable import AudioSnapshotTesting

@Test(
    .audioSnapshot(record: false, strategy: .waveform(width: 3000, height: 800)),
    arguments: ["sine", "triangle", "square", "sawtooth", "brown", "pink", "white"]
)
func fileWaveform(wave: String) async throws {
    await assertAudioSnapshot(
        of: try AVAudioPCMBuffer.read(wave: wave),
        named: "fileWaveform.\(wave)"
    )
}

@Test(
    .audioSnapshot(record: false, strategy: .waveform(width: 800, height: 400, strategy: .individualLines))
)
func fileWaveformMetronome() async throws {
    await assertAudioSnapshot(
        of: try AVAudioPCMBuffer.read(wave: "metronome"),
        named: "fileWaveformMetronome.metronome"
    )
}

@Test(
    .audioSnapshot(record: false, strategy: .waveform(width: 4000, height: 1000, strategy: .individualLines, mono: false))
)
func stereoFileWaveform() async throws {
    await assertAudioSnapshot(
        of: try AVAudioPCMBuffer.read(wave: "left-right-metronome"),
        named: "stereoFileWaveform"
    )
}

@Test(
    .audioSnapshot(record: false, strategy: .waveform(width: 4000, height: 1000))
)
func fileWaveformOverlay() async throws {
    let buffer1 = try AVAudioPCMBuffer.read(wave: "sine")
    let buffer2 = try AVAudioPCMBuffer.read(wave: "square")

    await assertAudioSnapshot(
        of: [buffer1, buffer2],
        named: "fileWaveformOverlay.square-over-sine"
    )
}

@available(iOS 16, macOS 13, *)
@Test(
    .audioSnapshot(record: false),
    arguments: ["1hz@32768", "1hz2hz@32768", "2hz@32768"]
)
func perfectSpectrum(wave: String) async throws {
    let buffer = try AVAudioPCMBuffer.read(wave: wave)
    // Override with custom window for perfect spectrum analysis
    await AudioSnapshotConfiguration.$current.withValue(
        AudioSnapshotTrait(record: false, strategy: .spectrum(
            width: 500,
            height: 200,
            window: .init(repeating: 1, count: Int(buffer.frameLength))
        ))
    ) {
        await assertAudioSnapshot(of: buffer, named: "perfectSpectrum.\(wave)")
    }
}

@available(iOS 16, macOS 13, *)
@Test(
    .audioSnapshot(record: false, strategy: .spectrum(width: 1500, height: 400)),
    arguments: ["1hz@44100", "white", "brown", "pink", "square", "triangle", "sawtooth"]
)
func windowedSpectrum(wave: String) async throws {
    let buffer = try AVAudioPCMBuffer.read(wave: wave)
    await assertAudioSnapshot(
        of: buffer,
        named: "windowedSpectrum.\(wave)"
    )
}

@available(iOS 16, macOS 13, *)
@Test(
    .audioSnapshot(record: false, strategy: .spectrum(width: 1500, height: 400, window: .init(repeating: 1, count: 32768)))
)
func spectrumSynthesised() async throws {
    let signal = synthesizeSignal(frequencyAmplitudePairs: (0..<1024).map { (Float($0), Float($0) / Float(1024)) }, count: 32768)
    let buffer = createBuffer(from: signal)

    await assertAudioSnapshot(
        of: buffer,
        named: "spectrumSynthesised.synthesized"
    )
}

@available(iOS 16, macOS 13, *)
@Test(
    "Generates color spectrum spectrogram",
    .audioSnapshot(
        record: false,
        strategy: .spectrogram(hopSize: 128, frequencyCount: 1024, window: [Float](repeating: 1, count: 2048), amplitudeScale: .linear, imageWidth: 200)
    )
)
func spectrogramColors() async throws {
    let frequencyCount = 1024
    let signal = synthesizeSignal(
        frequencyAmplitudePairs: (0..<frequencyCount).map { (Float($0), Float($0) / Float(frequencyCount)) },
        count: 2048
    )
    let buffer = createBuffer(from: signal + signal + signal + signal)

    await assertAudioSnapshot(of: buffer, named: "spectrogramColors")
}

@available(iOS 16, macOS 13, *)
@Test(
    .audioSnapshot(record: false, strategy: .spectrogram(hopSize: 4096, frequencyCount: 2048)),
    arguments: [
        "500hz",
        "2000hz",
        "5000hz",
        "15000hz",
        "20000hz",
        "chirp",
        "440-880-1320hz",
        "beating",
        "fade-in-out",
        "ultra-high-tones",
        "high-sweep",
        "spring"
    ]
)
func spectrogramHop256(wave: String) async throws {
    let buffer = try AVAudioPCMBuffer.read(wave: wave)
    await assertAudioSnapshot(
        of: buffer,
        named: "spectrogramHop256.\(wave)"
    )
}

@Test(
    "Audio snapshot of synthesized signal",
    .audioSnapshot(record: false, strategy: .waveform(width: 1000, height: 300))
)
func audioSnapshotSynthesized() async throws {
    let signal = synthesizeSignal(
        frequencyAmplitudePairs: [(440, 0.5), (880, 0.3)],
        count: 4410
    )
    let buffer = createBuffer(from: signal)

    await assertAudioSnapshot(
        of: buffer,
        named: "audioSnapshotSynthesized.440-880hz"
    )
}

@available(iOS 16, macOS 13, *)
@Test(
    "Multi-channel (4ch) audio comparison",
    .audioSnapshot(record: false, strategy: .spectrogram(hopSize: 4096, frequencyCount: 2048))
)
func multiChannelComparison() async throws {
    let frameCount = 44100
    let frequencies: [Float] = [440, 880, 1320, 1760]
    let channels = frequencies.map { frequency in
        synthesizeSignal(frequencyAmplitudePairs: [(frequency, 0.5)], count: frameCount)
    }
    let buffer = createBuffer(channels: channels, sampleRate: 44100)
    await assertAudioSnapshot(of: buffer, named: "multiChannelComparison.4ch")
}

@Test(
    "Checksum snapshot records and verifies a deterministic buffer",
    .audioSnapshot(record: false, format: .checksum)
)
func checksumRoundTrip() async throws {
    let signal = synthesizeSignal(
        frequencyAmplitudePairs: [(440, 0.5)],
        count: 4410
    )
    let buffer = createBuffer(from: signal)
    await assertAudioSnapshot(of: buffer, named: "checksumRoundTrip.440hz")
}

@Test(
    "Checksum snapshot with multiple buffers uses indexed naming",
    .audioSnapshot(record: false, format: .checksum)
)
func checksumMultiBuffer() async throws {
    let signal1 = synthesizeSignal(
        frequencyAmplitudePairs: [(440, 0.5)],
        count: 4410
    )
    let signal2 = synthesizeSignal(
        frequencyAmplitudePairs: [(880, 0.3)],
        count: 4410
    )
    let buffer1 = createBuffer(from: signal1)
    let buffer2 = createBuffer(from: signal2)
    await assertAudioSnapshot(
        of: [buffer1, buffer2],
        named: "checksumMultiBuffer"
    )
}

private func createBuffer(from samples: [Float], sampleRate: Double = 32768) -> AVAudioPCMBuffer {
    createBuffer(channels: [samples], sampleRate: sampleRate)
}

private func createBuffer(channels: [[Float]], sampleRate: Double = 32768) -> AVAudioPCMBuffer {
    let channelCount = channels.count
    let frameCount = channels.first?.count ?? 0

    let format: AVAudioFormat
    if channelCount <= 2 {
        format = AVAudioFormat(standardFormatWithSampleRate: sampleRate, channels: AVAudioChannelCount(channelCount))!
    } else {
        let layout = AVAudioChannelLayout(layoutTag: kAudioChannelLayoutTag_DiscreteInOrder | UInt32(channelCount))!
        format = AVAudioFormat(standardFormatWithSampleRate: sampleRate, channelLayout: layout)
    }

    let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: AVAudioFrameCount(frameCount))!
    for (channelIndex, samples) in channels.enumerated() {
        let channelData = buffer.floatChannelData![channelIndex]
        for (frameIndex, sample) in samples.enumerated() {
            channelData[frameIndex] = sample
        }
    }
    buffer.frameLength = AVAudioFrameCount(frameCount)
    return buffer
}

private extension AVAudioPCMBuffer {
    static func read(wave: String) throws -> AVAudioPCMBuffer {
        let file = try AVAudioFile(
            forReading: Bundle.module.url(forResource: wave, withExtension: "wav")!
        )
        let buffer = try #require(AVAudioPCMBuffer(pcmFormat: file.processingFormat, frameCapacity: AVAudioFrameCount(file.length)))
        try file.read(into: buffer)
        return buffer
    }
}

/// Synthesize a signal composed of multiple sine waves given frequency-amplitude pairs.
private func synthesizeSignal(
    frequencyAmplitudePairs: [(f: Float, a: Float)],
    count: Int
) -> [Float] {
    let tau: Float = .pi * 2
    let signal: [Float] = (0 ..< count).map { index in
        frequencyAmplitudePairs.reduce(0) { accumulator, frequenciesAmplitudePair in
            let normalizedIndex = Float(index) / Float(count)
            return accumulator + sin(normalizedIndex * frequenciesAmplitudePair.f * tau) * frequenciesAmplitudePair.a
        }
    }

    return signal
}