Video-Output/tests/decode.rs

//! End-to-end MJPEG decode tests against real JPEG files.
//!
//! Fixtures are solid green -- the colour the ATEM chroma keys on -- so a wrong
//! YUV matrix, a swapped Cb/Cr, or a studio-vs-full range mistake fails here
//! rather than on stage. Regenerate them with `python3 tests/data/generate.py`.

use video_output::frame::{NullPool, PixelFormat};
use video_output::inputs::mjpeg::MjpegDecoder;

const W: u32 = 64;
const H: u32 = 32;

/// Full-range BT.601 for pure green (0, 255, 0), which is what JFIF stores.
const GREEN_Y: u8 = 150;
const GREEN_CB: u8 = 44;
const GREEN_CR: u8 = 21;

/// JPEG at quality 95 on a flat colour lands within a couple of levels.
const TOL: i32 = 4;

fn fixture(name: &str) -> Vec<u8> {
    let p = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/");
    std::fs::read(format!("{p}{name}")).unwrap_or_else(|e| panic!("reading {name}: {e}"))
}

fn close(got: u8, want: u8, label: &str) {
    let d = (got as i32 - want as i32).abs();
    assert!(d <= TOL, "{label}: got {got}, expected ~{want} (tolerance {TOL})");
}

#[test]
fn decodes_422_into_packed_yuy2() {
    let mut d = MjpegDecoder::new().unwrap();
    let frame = d.decode(&fixture("green-422.jpg"), &NullPool, 1).unwrap();

    assert_eq!((frame.width, frame.height), (W, H));
    assert_eq!(frame.format, PixelFormat::Yuy2, "4:2:2 must take the zero-conversion path");
    assert_eq!(frame.layout.total, (W * H * 2) as usize);
    assert_eq!(frame.seq, 1);

    // YUY2 byte order is Y0 U Y1 V.
    let p = frame.plane(0);
    close(p[0], GREEN_Y, "Y0");
    close(p[1], GREEN_CB, "U");
    close(p[2], GREEN_Y, "Y1");
    close(p[3], GREEN_CR, "V");

    // The whole frame is one colour, so the pattern must hold everywhere --
    // this catches a stride or interleave bug that only shows on later rows.
    for (i, quad) in p.chunks_exact(4).enumerate() {
        close(quad[0], GREEN_Y, &format!("Y0 at quad {i}"));
        close(quad[1], GREEN_CB, &format!("U at quad {i}"));
        close(quad[2], GREEN_Y, &format!("Y1 at quad {i}"));
        close(quad[3], GREEN_CR, &format!("V at quad {i}"));
    }
}

#[test]
fn decodes_420_into_planar_i420() {
    let mut d = MjpegDecoder::new().unwrap();
    let frame = d.decode(&fixture("green-420.jpg"), &NullPool, 2).unwrap();

    assert_eq!(frame.format, PixelFormat::I420);
    assert_eq!(frame.layout.total, (W * H * 3 / 2) as usize);

    let y = frame.plane(0);
    let u = frame.plane(1);
    let v = frame.plane(2);
    assert_eq!(y.len(), (W * H) as usize);
    assert_eq!(u.len(), (W / 2 * H / 2) as usize);
    assert_eq!(v.len(), (W / 2 * H / 2) as usize);

    for (i, &s) in y.iter().enumerate() {
        close(s, GREEN_Y, &format!("Y[{i}]"));
    }
    for (i, &s) in u.iter().enumerate() {
        close(s, GREEN_CB, &format!("U[{i}]"));
    }
    for (i, &s) in v.iter().enumerate() {
        close(s, GREEN_CR, &format!("V[{i}]"));
    }
}

#[test]
fn decodes_444_through_the_bgrx_fallback() {
    let mut d = MjpegDecoder::new().unwrap();
    let frame = d.decode(&fixture("green-444.jpg"), &NullPool, 3).unwrap();

    assert_eq!(frame.format, PixelFormat::Bgrx, "4:4:4 has no packed YUV equivalent");
    assert_eq!(frame.layout.total, (W * H * 4) as usize);

    // Byte order is B, G, R, X.
    for (i, px) in frame.plane(0).chunks_exact(4).enumerate() {
        close(px[0], 0, &format!("B at {i}"));
        close(px[1], 255, &format!("G at {i}"));
        close(px[2], 0, &format!("R at {i}"));
    }
}

#[test]
fn decodes_grayscale_through_the_bgrx_fallback() {
    let mut d = MjpegDecoder::new().unwrap();
    let frame = d.decode(&fixture("gray.jpg"), &NullPool, 4).unwrap();

    assert_eq!(frame.format, PixelFormat::Bgrx);
    for px in frame.plane(0).chunks_exact(4) {
        close(px[0], 128, "B");
        close(px[1], 128, "G");
        close(px[2], 128, "R");
    }
}

#[test]
fn header_read_reports_dimensions_and_subsampling() {
    let mut d = MjpegDecoder::new().unwrap();
    let info = d.read_header(&fixture("green-422.jpg")).unwrap();
    assert_eq!((info.width, info.height), (W, H));
    assert_eq!(info.subsamp_name(), "4:2:2");
    assert_eq!(info.colorspace_name(), "YCbCr");
    assert_eq!(info.target_format(), PixelFormat::Yuy2);
}

/// The failure mode that actually happens on USB: a frame that stops early.
/// Every truncation length must be rejected without panicking or reading out
/// of bounds.
#[test]
fn every_truncation_of_a_real_jpeg_fails_cleanly() {
    let full = fixture("green-422.jpg");
    let mut d = MjpegDecoder::new().unwrap();

    for cut in 1..full.len() {
        let partial = &full[..cut];
        // Must not panic. Success is acceptable only at the full length.
        let _ = d.decode(partial, &NullPool, cut as u64);
    }

    // The decoder must still work afterwards -- no poisoned state.
    let frame = d.decode(&full, &NullPool, 999).unwrap();
    assert_eq!(frame.format, PixelFormat::Yuy2);
    assert_eq!(frame.seq, 999);
}

/// Bit flips inside the entropy-coded data are the other common USB fault.
#[test]
fn corrupted_payloads_never_panic_and_leave_the_decoder_usable() {
    let full = fixture("green-422.jpg");
    let mut d = MjpegDecoder::new().unwrap();

    for seed in 0..200u32 {
        let mut bad = full.clone();
        let i = (seed as usize * 7 + 100) % bad.len();
        bad[i] ^= 0xFF;
        let _ = d.decode(&bad, &NullPool, seed as u64);
    }

    assert!(d.decode(&full, &NullPool, 1).is_ok(), "decoder must survive a burst of corruption");
}

#[test]
fn non_jpeg_input_is_rejected() {
    let mut d = MjpegDecoder::new().unwrap();
    assert!(d.decode(b"not a jpeg at all, not even close", &NullPool, 0).is_err());
    assert!(d.decode(&[], &NullPool, 0).is_err());
    assert!(d.decode(&[0xFF; 1024], &NullPool, 0).is_err());
}

/// Decoding many frames through one decoder must not grow memory or drift --
/// this is the steady state the appliance actually runs in.
#[test]
fn repeated_decodes_are_stable() {
    let data = fixture("green-422.jpg");
    let mut d = MjpegDecoder::new().unwrap();
    for seq in 0..100 {
        let f = d.decode(&data, &NullPool, seq).unwrap();
        assert_eq!(f.layout.total, (W * H * 2) as usize);
        close(f.plane(0)[0], GREEN_Y, "Y across repeats");
    }
}