tests: Some more tests

Signed-off-by: Alexandre Milesi <[email protected]>

Author: milesial

Cargo.lock

@@ -172,6 +172,7 @@ insta = { version = "1.41", features = [
 ] }
 aligned-vec = "0.6.4"
 lazy_static = "1.4"
+mockito = "1.7.0"
 
 [build-dependencies]
 tonic-build = { version = "0.13.1" }

lib/llm/Cargo.toml

@@ -16,6 +16,7 @@ use crate::preprocessor::media::{ImageDecoder, VideoDecoder};
 use nixl_sys::Agent as NixlAgent;
 
 // Raw encoded media data (.png, .mp4, ...), optionally b64-encoded
+#[derive(Debug)]
 pub struct EncodedMediaData {
     pub(crate) bytes: Vec<u8>,
     pub(crate) b64_encoded: bool,
@@ -198,3 +199,97 @@ impl MediaLoader {
         Ok(rdma_descriptor)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_from_base64() {
+        // Simple base64 encoded "test" string: dGVzdA==
+        let data_url = url::Url::parse("data:text/plain;base64,dGVzdA==").unwrap();
+        let client = reqwest::Client::new();
+
+        let result = EncodedMediaData::from_url(&data_url, &client)
+            .await
+            .unwrap();
+
+        assert!(result.b64_encoded);
+        assert_eq!(result.bytes, b"dGVzdA==");
+        let decoded = result.into_bytes().unwrap();
+        assert_eq!(decoded, b"test");
+    }
+
+    #[tokio::test]
+    async fn test_from_empty_base64() {
+        let data_url = url::Url::parse("data:text/plain;base64,").unwrap();
+        let client = reqwest::Client::new();
+
+        let result = EncodedMediaData::from_url(&data_url, &client).await;
+        assert!(result.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_from_invalid_base64() {
+        let data_url = url::Url::parse("data:invalid").unwrap();
+        let client = reqwest::Client::new();
+
+        let result = EncodedMediaData::from_url(&data_url, &client).await;
+        assert!(result.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_from_url_http() {
+        let mut server = mockito::Server::new_async().await;
+        let mock = server
+            .mock("GET", "/image.png")
+            .with_status(200)
+            .with_body(b"test data")
+            .create_async()
+            .await;
+
+        let url = url::Url::parse(&format!("{}/image.png", server.url())).unwrap();
+        let client = reqwest::Client::new();
+
+        let result = EncodedMediaData::from_url(&url, &client).await.unwrap();
+
+        assert!(!result.b64_encoded);
+        assert_eq!(result.bytes, b"test data");
+        let decoded = result.into_bytes().unwrap();
+        assert_eq!(decoded, b"test data");
+
+        mock.assert_async().await;
+    }
+
+    #[tokio::test]
+    async fn test_from_url_http_404() {
+        let mut server = mockito::Server::new_async().await;
+        let mock = server
+            .mock("GET", "/image.png")
+            .with_status(404)
+            .create_async()
+            .await;
+
+        let url = url::Url::parse(&format!("{}/image.png", server.url())).unwrap();
+        let client = reqwest::Client::new();
+        let result = EncodedMediaData::from_url(&url, &client).await;
+        assert!(result.is_err());
+
+        mock.assert_async().await;
+    }
+
+    #[tokio::test]
+    async fn test_from_unsupported_scheme() {
+        let ftp_url = url::Url::parse("ftp://example.com/image.png").unwrap();
+        let client = reqwest::Client::new();
+
+        let result = EncodedMediaData::from_url(&ftp_url, &client).await;
+        assert!(result.is_err());
+        assert!(
+            result
+                .unwrap_err()
+                .to_string()
+                .contains("Unsupported media URL scheme")
+        );
+    }
+}

lib/llm/src/preprocessor/media/common.rs

@@ -89,25 +89,15 @@ mod tests {
     #[case(1, image::ImageFormat::Png, 8, 12, 1, "Grayscale PNG")]
     #[case(3, image::ImageFormat::Jpeg, 15, 20, 3, "RGB JPEG")]
     #[case(3, image::ImageFormat::Bmp, 12, 18, 3, "RGB BMP")]
-    #[case(4, image::ImageFormat::Bmp, 7, 9, 4, "RGBA BMP")]
-    #[case(1, image::ImageFormat::Bmp, 5, 6, 3, "Grayscale BMP")] // BMP converts grayscale to RGB
-    #[case(3, image::ImageFormat::Gif, 10, 10, 4, "RGB GIF")] // GIF may add alpha channel
     #[case(3, image::ImageFormat::WebP, 8, 8, 3, "RGB WebP")]
-    #[case(4, image::ImageFormat::WebP, 9, 11, 4, "RGBA WebP")]
-    #[case(1, image::ImageFormat::WebP, 6, 7, 3, "Grayscale WebP")] // WebP converts grayscale to RGB
-    fn test_decode_image_formats(
+    fn test_image_decode(
         #[case] input_channels: u32,
         #[case] format: image::ImageFormat,
         #[case] width: u32,
         #[case] height: u32,
         #[case] expected_channels: u32,
         #[case] description: &str,
     ) {
-        // Skip JPEG for non-RGB formats (JPEG doesn't support transparency or pure grayscale)
-        if format == image::ImageFormat::Jpeg && input_channels != 3 {
-            return;
-        }
-
         let decoder = ImageDecoder::default();
         let image_bytes = create_test_image(width, height, input_channels, format);
         let encoded_data = create_encoded_media_data(image_bytes);
@@ -124,11 +114,9 @@ mod tests {
     }
 
     #[rstest]
-    #[case(Some(100), 8, 10, image::ImageFormat::Png, true, "within limit")] // 80 pixels < 100
-    #[case(Some(50), 10, 10, image::ImageFormat::Jpeg, false, "exceeds limit")] // 100 pixels > 50
-    #[case(Some(25), 5, 5, image::ImageFormat::Bmp, true, "exactly at limit")] // 25 pixels = 25
-    #[case(None, 200, 300, image::ImageFormat::Png, true, "no limit")] // 60,000 pixels, no limit
-    #[case(Some(100), 9, 10, image::ImageFormat::WebP, true, "webp within limit")] // 90 pixels < 100
+    #[case(Some(200), 10, 10, image::ImageFormat::Png, true, "within limit")]
+    #[case(Some(50), 10, 10, image::ImageFormat::Jpeg, false, "exceeds limit")]
+    #[case(None, 200, 300, image::ImageFormat::Png, true, "no limit")]
     fn test_pixel_limits(
         #[case] max_pixels: Option<usize>,
         #[case] width: u32,
@@ -173,44 +161,10 @@ mod tests {
         }
     }
 
-    #[test]
-    fn test_invalid_image_data() {
-        let decoder = ImageDecoder::default();
-        // Random bytes that are not a valid image
-        let invalid_bytes = vec![0xFF, 0xFE, 0xFD, 0xFC, 0xFB, 0xFA, 0xF9, 0xF8];
-        let encoded_data = create_encoded_media_data(invalid_bytes);
-
-        let result = decoder.decode(encoded_data);
-        assert!(
-            result.is_err(),
-            "Should fail when decoding invalid image data"
-        );
-    }
-
-    #[test]
-    fn test_empty_image_data() {
-        let decoder = ImageDecoder::default();
-        let empty_bytes = vec![];
-        let encoded_data = create_encoded_media_data(empty_bytes);
-
-        let result = decoder.decode(encoded_data);
-        assert!(result.is_err(), "Should fail when decoding empty data");
-    }
-
     #[rstest]
     #[case(3, image::ImageFormat::Png)]
-    #[case(4, image::ImageFormat::Png)]
-    #[case(1, image::ImageFormat::Png)]
-    #[case(3, image::ImageFormat::Bmp)]
-    #[case(1, image::ImageFormat::Bmp)]
-    #[case(3, image::ImageFormat::Jpeg)]
-    #[case(3, image::ImageFormat::WebP)]
-    #[case(4, image::ImageFormat::WebP)]
-    #[case(1, image::ImageFormat::WebP)]
-    #[case(3, image::ImageFormat::Gif)]
-    fn test_small_edge_case(#[case] input_channels: u32, #[case] format: image::ImageFormat) {
+    fn test_decode_1x1_image(#[case] input_channels: u32, #[case] format: image::ImageFormat) {
         let decoder = ImageDecoder::default();
-        // Test with 1x1 image (smallest possible)
         let image_bytes = create_test_image(1, 1, input_channels, format);
         let encoded_data = create_encoded_media_data(image_bytes);
 

lib/llm/src/preprocessor/media/image.rs

@@ -82,15 +82,13 @@ impl Decoder for VideoDecoder {
         let mut num_frames_decoded = 0;
 
         let target_indices = if requested_frames == 1 {
-            vec![(total_frames / 2)]
+            vec![total_frames / 2]
         } else {
             (0..requested_frames)
                 .map(|i| (i * (total_frames - 1)) / (requested_frames - 1))
                 .collect()
         };
 
-        println!("target_indices: {:?}", target_indices);
-
         // Decode all frames sequentially (required for P/B-frames), but only keep target frames
         // TODO: smarter seek-based decoding for better sparse sampling
         for (current_frame_idx, result) in decoder.decode_iter().enumerate() {
@@ -363,31 +361,6 @@ mod tests {
         }
     }
 
-    // Invalid/Edge Cases
-
-    #[test]
-    fn test_invalid_video_data() {
-        let decoder = VideoDecoder::default();
-        let invalid_bytes = vec![0xFF, 0xFE, 0xFD, 0xFC, 0xFB, 0xFA, 0xF9, 0xF8];
-        let encoded_data = create_encoded_media_data(invalid_bytes);
-
-        let result = decoder.decode(encoded_data);
-        assert!(
-            result.is_err(),
-            "Should fail when decoding invalid video data"
-        );
-    }
-
-    #[test]
-    fn test_empty_video_data() {
-        let decoder = VideoDecoder::default();
-        let empty_bytes = vec![];
-        let encoded_data = create_encoded_media_data(empty_bytes);
-
-        let result = decoder.decode(encoded_data);
-        assert!(result.is_err(), "Should fail when decoding empty data");
-    }
-
     #[test]
     fn test_conflicting_fps_and_num_frames() {
         let video_bytes = load_test_video("240p_10.mp4");
@@ -410,85 +383,4 @@ mod tests {
         let error_msg = result.unwrap_err().to_string();
         assert!(error_msg.contains("cannot be specified at the same time"));
     }
-
-    #[test]
-    fn test_conflicting_max_frames_and_num_frames() {
-        let video_bytes = load_test_video("240p_10.mp4");
-
-        let decoder = VideoDecoder {
-            fps: None,
-            max_frames: Some(3),
-            num_frames: Some(5),
-            strict: false,
-            max_pixels: None,
-        };
-
-        let encoded_data = create_encoded_media_data(video_bytes);
-
-        let result = decoder.decode(encoded_data);
-        assert!(
-            result.is_err(),
-            "Should fail when both max_frames and num_frames are specified"
-        );
-        let error_msg = result.unwrap_err().to_string();
-        assert!(error_msg.contains("cannot be specified at the same time"));
-    }
-
-    #[test]
-    fn test_strict_mode_success() {
-        let video_bytes = load_test_video("240p_10.mp4");
-
-        let decoder = VideoDecoder {
-            fps: None,
-            max_frames: None,
-            num_frames: Some(5),
-            strict: true,
-            max_pixels: None,
-        };
-
-        let encoded_data = create_encoded_media_data(video_bytes);
-
-        let result = decoder.decode(encoded_data);
-        assert!(
-            result.is_ok(),
-            "Should succeed in strict mode when all frames decode successfully"
-        );
-
-        let decoded = result.unwrap();
-        assert_eq!(
-            decoded.shape[0], 5,
-            "Should decode exactly 5 frames in strict mode"
-        );
-    }
-
-    #[test]
-    fn test_small_video_edge_case() {
-        let video_bytes = load_test_video("2p_10.mp4");
-        let (width, height, expected_frames) = parse_video_info("2p_10.mp4");
-
-        let decoder = VideoDecoder::default();
-        let encoded_data = create_encoded_media_data(video_bytes);
-
-        let result = decoder.decode(encoded_data);
-        assert!(result.is_ok(), "Should decode 2x2 small video successfully");
-
-        let decoded = result.unwrap();
-        assert_eq!(decoded.shape.len(), 4, "Should have 4 dimensions");
-        assert_eq!(
-            decoded.shape[0], expected_frames as usize,
-            "Should decode all frames"
-        );
-        assert_eq!(
-            decoded.shape[1], height as usize,
-            "Height should be {}",
-            height
-        );
-        assert_eq!(
-            decoded.shape[2], width as usize,
-            "Width should be {}",
-            width
-        );
-        assert_eq!(decoded.shape[3], 3, "Channels should be 3");
-        assert_eq!(decoded.dtype, "uint8");
-    }
 }