[tests] Add tests that verify YUV conversion (#860)

src/snapshot_tests.rs

@@ -24,6 +24,7 @@ mod tiles_tests;
 mod tiles_transitions_tests;
 mod transition_tests;
 mod view_tests;
+mod yuv_tests;
 
 const DEFAULT_RESOLUTION: Resolution = Resolution {
     width: 640,

src/snapshot_tests/test_case.rs

@@ -25,6 +25,7 @@ pub(super) struct TestCase {
     pub only: bool,
     pub allowed_error: f32,
     pub resolution: Resolution,
+    pub output_format: OutputFrameFormat,
 }
 
 impl Default for TestCase {
@@ -41,6 +42,7 @@ impl Default for TestCase {
                 width: 640,
                 height: 360,
             },
+            output_format: OutputFrameFormat::PlanarYuv420Bytes,
         }
     }
 }
@@ -51,8 +53,8 @@ pub(super) enum TestResult {
 }
 
 impl TestCase {
-    fn renderer(&self) -> Renderer {
-        let renderer = create_renderer();
+    pub(super) fn renderer(&self) -> Renderer {
+        let mut renderer = create_renderer();
         for (id, spec) in self.renderers.iter() {
             renderer
                 .register_renderer(id.clone(), spec.clone())
@@ -63,27 +65,27 @@ impl TestCase {
             renderer.register_input(InputId(format!("input_{}", index + 1).into()))
         }
 
-        renderer
-    }
-
-    pub(super) fn run(&self) -> TestResult {
-        if self.name.is_empty() {
-            panic!("Snapshot test name has to be provided");
-        }
-        let mut renderer = self.renderer();
-        let mut result = TestResult::Success;
-
         for update in &self.scene_updates {
             renderer
                 .update_scene(
                     OutputId(OUTPUT_ID.into()),
                     self.resolution,
-                    OutputFrameFormat::PlanarYuv420Bytes,
+                    self.output_format,
                     update.clone(),
                 )
                 .unwrap();
         }
 
+        renderer
+    }
+
+    pub(super) fn run(&self) -> TestResult {
+        if self.name.is_empty() {
+            panic!("Snapshot test name has to be provided");
+        }
+        let mut renderer = self.renderer();
+        let mut result = TestResult::Success;
+
         for pts in self.timestamps.iter().copied() {
             if let TestResult::Failure = self.test_snapshots_for_pts(&mut renderer, pts) {
                 result = TestResult::Failure;
@@ -122,7 +124,11 @@ impl TestCase {
             .collect()
     }
 
-    fn snapshot_for_pts(&self, renderer: &mut Renderer, pts: Duration) -> Result<Snapshot> {
+    pub(super) fn snapshot_for_pts(
+        &self,
+        renderer: &mut Renderer,
+        pts: Duration,
+    ) -> Result<Snapshot> {
         let mut frame_set = FrameSet::new(pts);
         for input in self.inputs.iter() {
             let input_id = InputId::from(Arc::from(input.name.clone()));

src/snapshot_tests/utils.rs

@@ -1,25 +1,36 @@
 use core::panic;
 use std::{
+    io::Write,
     sync::{Arc, OnceLock},
     time::Duration,
 };
 
+use bytes::BufMut;
 use compositor_render::{
     create_wgpu_ctx, web_renderer, Frame, FrameData, Framerate, Renderer, RendererOptions,
     WgpuFeatures, YuvPlanes,
 };
+use crossbeam_channel::bounded;
+use tracing::error;
 
 pub const SNAPSHOTS_DIR_NAME: &str = "snapshot_tests/snapshots/render_snapshots";
 
 pub(super) fn frame_to_rgba(frame: &Frame) -> Vec<u8> {
-    let FrameData::PlanarYuv420(YuvPlanes {
+    match &frame.data {
+        FrameData::PlanarYuv420(planes) => yuv_frame_to_rgba(frame, planes),
+        FrameData::PlanarYuvJ420(_) => panic!("unsupported"),
+        FrameData::InterleavedYuv422(_) => panic!("unsupported"),
+        FrameData::Rgba8UnormWgpuTexture(texture) => read_rgba_texture(texture).to_vec(),
+        FrameData::Nv12WgpuTexture(_) => panic!("unsupported"),
+    }
+}
+
+pub(super) fn yuv_frame_to_rgba(frame: &Frame, planes: &YuvPlanes) -> Vec<u8> {
+    let YuvPlanes {
         y_plane,
         u_plane,
         v_plane,
-    }) = &frame.data
-    else {
-        panic!("Wrong pixel format")
-    };
+    } = planes;
 
     // Renderer can sometimes produce resolution that is not dividable by 2
     let corrected_width = frame.resolution.width - (frame.resolution.width % 2);
@@ -46,11 +57,13 @@ pub(super) fn frame_to_rgba(frame: &Frame) -> Vec<u8> {
     rgba_data
 }
 
-pub(super) fn create_renderer() -> Renderer {
+fn get_wgpu_ctx() -> (Arc<wgpu::Device>, Arc<wgpu::Queue>) {
     static CTX: OnceLock<(Arc<wgpu::Device>, Arc<wgpu::Queue>)> = OnceLock::new();
-    let wgpu_ctx =
-        CTX.get_or_init(|| create_wgpu_ctx(false, Default::default(), Default::default()).unwrap());
+    CTX.get_or_init(|| create_wgpu_ctx(false, Default::default(), Default::default()).unwrap())
+        .clone()
+}
 
+pub(super) fn create_renderer() -> Renderer {
     let (renderer, _event_loop) = Renderer::new(RendererOptions {
         web_renderer: web_renderer::WebRendererInitOptions {
             enable: false,
@@ -60,9 +73,98 @@ pub(super) fn create_renderer() -> Renderer {
         framerate: Framerate { num: 30, den: 1 },
         stream_fallback_timeout: Duration::from_secs(3),
         wgpu_features: WgpuFeatures::default(),
-        wgpu_ctx: Some(wgpu_ctx.clone()),
+        wgpu_ctx: Some(get_wgpu_ctx()),
         load_system_fonts: false,
     })
     .unwrap();
     renderer
 }
+
+fn read_rgba_texture(texture: &wgpu::Texture) -> bytes::Bytes {
+    let (device, queue) = get_wgpu_ctx();
+    let buffer = new_download_buffer(&device, texture);
+
+    let mut encoder = device.create_command_encoder(&Default::default());
+    copy_to_buffer(&mut encoder, texture, &buffer);
+    queue.submit(Some(encoder.finish()));
+
+    download_buffer(&device, texture.size(), &buffer)
+}
+
+fn new_download_buffer(device: &wgpu::Device, texture: &wgpu::Texture) -> wgpu::Buffer {
+    let size = texture.size();
+    let block_size = texture.format().block_copy_size(None).unwrap();
+
+    device.create_buffer(&wgpu::BufferDescriptor {
+        label: Some("texture buffer"),
+        mapped_at_creation: false,
+        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
+        size: (pad_to_256(block_size * size.width) * size.height) as u64,
+    })
+}
+
+fn copy_to_buffer(
+    encoder: &mut wgpu::CommandEncoder,
+    texture: &wgpu::Texture,
+    buffer: &wgpu::Buffer,
+) {
+    let size = texture.size();
+    let block_size = texture.format().block_copy_size(None).unwrap();
+    encoder.copy_texture_to_buffer(
+        wgpu::ImageCopyTexture {
+            aspect: wgpu::TextureAspect::All,
+            mip_level: 0,
+            origin: wgpu::Origin3d::ZERO,
+            texture,
+        },
+        wgpu::ImageCopyBuffer {
+            buffer,
+            layout: wgpu::ImageDataLayout {
+                bytes_per_row: Some(pad_to_256(size.width * block_size)),
+                rows_per_image: Some(size.height),
+                offset: 0,
+            },
+        },
+        size,
+    );
+}
+
+fn download_buffer(
+    device: &wgpu::Device,
+    size: wgpu::Extent3d,
+    source: &wgpu::Buffer,
+) -> bytes::Bytes {
+    let buffer = bytes::BytesMut::with_capacity((size.width * size.height * 4) as usize);
+    let (s, r) = bounded(1);
+    source
+        .slice(..)
+        .map_async(wgpu::MapMode::Read, move |result| {
+            if let Err(err) = s.send(result) {
+                error!("channel send error: {err}")
+            }
+        });
+
+    device.poll(wgpu::MaintainBase::Wait);
+
+    r.recv().unwrap().unwrap();
+    let mut buffer = buffer.writer();
+    {
+        let range = source.slice(..).get_mapped_range();
+        let chunks = range.chunks(pad_to_256(size.width * 4) as usize);
+        for chunk in chunks {
+            buffer
+                .write_all(&chunk[..(size.width * 4) as usize])
+                .unwrap();
+        }
+    };
+    source.unmap();
+    buffer.into_inner().into()
+}
+
+fn pad_to_256(value: u32) -> u32 {
+    if value % 256 == 0 {
+        value
+    } else {
+        value + (256 - (value % 256))
+    }
+}

src/snapshot_tests/yuv_tests.rs

@@ -0,0 +1,125 @@
+use core::panic;
+use std::{sync::Arc, time::Duration};
+
+use compositor_render::{
+    scene::{
+        BorderRadius, Component, Overflow, Position, RGBAColor, ShaderComponent, Size,
+        ViewChildrenDirection, ViewComponent,
+    },
+    shader::ShaderSpec,
+    OutputFrameFormat, RendererId, RendererSpec, Resolution,
+};
+
+use super::test_case::TestCase;
+
+fn run_case(test_case: TestCase, expected: &[u8]) {
+    let mut renderer = test_case.renderer();
+    let snapshot = test_case
+        .snapshot_for_pts(&mut renderer, Duration::ZERO)
+        .unwrap();
+    let failed = snapshot
+        .data
+        .iter()
+        .zip(expected)
+        .any(|(actual, expected)| u8::abs_diff(*actual, *expected) > 2);
+    if failed {
+        panic!("Sample mismatched {:?}", snapshot.data)
+    }
+}
+
+/// Test how yuv output is generated for smooth color change
+#[test]
+fn yuv_test_gradient() {
+    let shader_id = RendererId(Arc::from("example_shader"));
+    let width = 8;
+    let height = 2;
+
+    let yuv_case = TestCase {
+        scene_updates: vec![Component::Shader(ShaderComponent {
+            id: None,
+            children: vec![],
+            shader_id: shader_id.clone(),
+            shader_param: None,
+            size: Size {
+                width: width as f32,
+                height: height as f32,
+            },
+        })],
+        renderers: vec![(
+            shader_id.clone(),
+            RendererSpec::Shader(ShaderSpec {
+                source: include_str!("./yuv_tests/gradient.wgsl").into(),
+            }),
+        )],
+        resolution: Resolution { width, height },
+        ..Default::default()
+    };
+    let rgb_case = TestCase {
+        output_format: OutputFrameFormat::RgbaWgpuTexture,
+        ..yuv_case.clone()
+    };
+
+    #[rustfmt::skip]
+    run_case(
+        yuv_case,
+        &[
+            91, 0, 0, 255, 106, 6, 5, 255, 161, 0, 0, 255, 169, 3, 3, 255, 204, 0, 0, 255, 210, 2, 2, 255, 238, 0, 0, 255, 242, 2, 1, 255,
+            91, 0, 0, 255, 106, 6, 5, 255, 161, 0, 0, 255, 169, 3, 3, 255, 204, 0, 0, 255, 210, 2, 2, 255, 238, 0, 0, 255, 242, 2, 1, 255,
+        ],
+    );
+    #[rustfmt::skip]
+    run_case(rgb_case,
+        &[
+            71, 0, 0, 255, 120, 0, 0, 255, 152, 0, 0, 255, 177, 0, 0, 255, 198, 0, 0, 255, 216, 0, 0, 255, 233, 0, 0, 255, 248, 0, 0, 255,
+            71, 0, 0, 255, 120, 0, 0, 255, 152, 0, 0, 255, 177, 0, 0, 255, 198, 0, 0, 255, 216, 0, 0, 255, 233, 0, 0, 255, 248, 0, 0, 255,
+        ],
+    );
+}
+
+/// Test how yuv output is generated for unified color
+#[test]
+fn yuv_test_uniform_color() {
+    let width = 8;
+    let height = 2;
+
+    let yuv_case = TestCase {
+        scene_updates: vec![Component::View(ViewComponent {
+            id: None,
+            children: vec![],
+            direction: ViewChildrenDirection::Row,
+            position: Position::Static {
+                width: None,
+                height: None,
+            },
+            transition: None,
+            overflow: Overflow::Hidden,
+            background_color: RGBAColor(50, 0, 0, 255),
+            border_radius: BorderRadius::ZERO,
+            border_width: 0.0,
+            border_color: RGBAColor(0, 0, 0, 0),
+            box_shadow: vec![],
+        })],
+        resolution: Resolution { width, height },
+        ..Default::default()
+    };
+    let rgb_case = TestCase {
+        output_format: OutputFrameFormat::RgbaWgpuTexture,
+        ..yuv_case.clone()
+    };
+
+    #[rustfmt::skip]
+    run_case(
+        yuv_case,
+        &[
+            50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255,
+            50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255
+        ],
+    );
+    #[rustfmt::skip]
+    run_case(rgb_case,
+        &[
+            50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255,
+            50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255, 50, 0, 0, 255
+        ],
+    );
+}