Skip to main content

immich_lib/testing/
generator.rs

1//! Test image generator for integration tests.
2//!
3//! Creates test images by transforming real base photos, ensuring
4//! CLIP-based duplicate detection works correctly in Immich.
5
6use std::path::{Path, PathBuf};
7use std::process::Command;
8
9use chrono::{DateTime, Utc};
10
11use crate::error::{ImmichError, Result};
12
13/// Transform specification for creating image variants.
14///
15/// Specifies how to transform a base image to create a test fixture.
16/// All fixtures in a duplicate group should use the same base image
17/// with different transforms (size, quality) to ensure CLIP sees them
18/// as duplicates.
19#[derive(Debug, Clone)]
20pub struct TransformSpec {
21    /// Base image filename (in tests/fixtures/base/)
22    pub base_image: String,
23    /// Target width in pixels (None = use base image width)
24    pub width: Option<u32>,
25    /// Target height in pixels (None = scale proportionally from width)
26    pub height: Option<u32>,
27    /// JPEG quality 1-100 (default 85)
28    pub quality: u8,
29    /// Strip dimension EXIF tags (for testing missing dimensions)
30    pub strip_dimensions: bool,
31}
32
33impl TransformSpec {
34    /// Create a new transform spec with default quality.
35    pub fn new(base_image: impl Into<String>) -> Self {
36        Self {
37            base_image: base_image.into(),
38            width: None,
39            height: None,
40            quality: 85,
41            strip_dimensions: false,
42        }
43    }
44
45    /// Set target dimensions.
46    pub fn with_size(mut self, width: u32, height: u32) -> Self {
47        self.width = Some(width);
48        self.height = Some(height);
49        self
50    }
51
52    /// Scale to a percentage of original size.
53    pub fn with_scale(mut self, scale_percent: u32) -> Self {
54        // Will be applied during generation by reading base image dimensions
55        // Store as negative width to signal percentage scaling
56        self.width = Some(scale_percent);
57        self.height = None; // Signal proportional scaling
58        self
59    }
60
61    /// Set JPEG quality.
62    pub fn with_quality(mut self, quality: u8) -> Self {
63        self.quality = quality;
64        self
65    }
66
67    /// Strip dimension EXIF tags from output.
68    pub fn without_dimensions(mut self) -> Self {
69        self.strip_dimensions = true;
70        self
71    }
72}
73
74impl Default for TransformSpec {
75    fn default() -> Self {
76        Self::new("base_landscape.jpg")
77    }
78}
79
80/// EXIF metadata specification.
81#[derive(Debug, Clone, Default)]
82pub struct ExifSpec {
83    /// GPS coordinates (latitude, longitude)
84    pub gps: Option<(f64, f64)>,
85    /// Capture datetime
86    pub datetime: Option<DateTime<Utc>>,
87    /// Timezone string (e.g., "+05:00")
88    pub timezone: Option<String>,
89    /// Camera manufacturer
90    pub camera_make: Option<String>,
91    /// Camera model
92    pub camera_model: Option<String>,
93    /// Image description
94    pub description: Option<String>,
95}
96
97/// Complete test image specification.
98#[derive(Debug, Clone)]
99pub struct TestImage {
100    /// Output filename
101    pub filename: String,
102    /// Transform to apply to base image
103    pub transform: TransformSpec,
104    /// EXIF metadata to embed
105    pub exif: ExifSpec,
106}
107
108impl TestImage {
109    /// Create a new test image specification.
110    pub fn new(filename: impl Into<String>, transform: TransformSpec) -> Self {
111        Self {
112            filename: filename.into(),
113            transform,
114            exif: ExifSpec::default(),
115        }
116    }
117
118    /// Add EXIF metadata.
119    pub fn with_exif(mut self, exif: ExifSpec) -> Self {
120        self.exif = exif;
121        self
122    }
123}
124
125/// Generate a test image by transforming a base image.
126///
127/// Loads a base image from `base_dir`, applies transforms (resize, recompress),
128/// saves to `output_dir`, and applies EXIF metadata.
129///
130/// # Arguments
131/// * `spec` - The test image specification
132/// * `base_dir` - Directory containing base images
133/// * `output_dir` - Directory to write the output image
134///
135/// # Returns
136/// Path to the generated image file
137pub fn generate_image(spec: &TestImage, base_dir: &Path, output_dir: &Path) -> Result<PathBuf> {
138    use image::imageops::FilterType;
139    use image::ImageFormat;
140
141    let ext = Path::new(&spec.filename)
142        .extension()
143        .and_then(|e| e.to_str())
144        .unwrap_or("")
145        .to_lowercase();
146
147    let output_path = output_dir.join(&spec.filename);
148
149    // Handle special formats
150    match ext.as_str() {
151        "mp4" | "mov" | "avi" => {
152            return generate_video(&spec.filename, output_dir, spec.transform.width, spec.transform.height);
153        }
154        "heic" | "heif" => {
155            return Err(ImmichError::Io(std::io::Error::other(
156                "HEIC encoding not available - requires platform-specific encoder",
157            )));
158        }
159        "cr3" | "cr2" | "nef" | "arw" | "dng" | "raf" | "orf" => {
160            return Err(ImmichError::Io(std::io::Error::other(
161                format!("RAW format .{} encoding not available - requires proprietary encoder", ext),
162            )));
163        }
164        _ => {}
165    }
166
167    // Load base image
168    let base_path = base_dir.join(&spec.transform.base_image);
169    let img = image::open(&base_path).map_err(|e| {
170        ImmichError::Io(std::io::Error::other(format!(
171            "Failed to load base image {}: {}",
172            base_path.display(),
173            e
174        )))
175    })?;
176
177    // Calculate target dimensions
178    let (target_width, target_height) = match (spec.transform.width, spec.transform.height) {
179        (Some(w), Some(h)) => (w, h),
180        (Some(scale), None) if scale <= 100 => {
181            // Interpret as percentage scale
182            let w = (img.width() * scale) / 100;
183            let h = (img.height() * scale) / 100;
184            (w.max(1), h.max(1))
185        }
186        (Some(w), None) => {
187            // Scale height proportionally
188            let h = (img.height() * w) / img.width();
189            (w, h.max(1))
190        }
191        (None, Some(h)) => {
192            // Scale width proportionally
193            let w = (img.width() * h) / img.height();
194            (w.max(1), h)
195        }
196        (None, None) => (img.width(), img.height()),
197    };
198
199    // Resize if needed
200    let resized = if target_width != img.width() || target_height != img.height() {
201        img.resize_exact(target_width, target_height, FilterType::Lanczos3)
202    } else {
203        img
204    };
205
206    // Save with specified quality
207    match ext.as_str() {
208        "png" => {
209            resized
210                .save_with_format(&output_path, ImageFormat::Png)
211                .map_err(|e| {
212                    ImmichError::Io(std::io::Error::other(format!("Failed to save PNG: {}", e)))
213                })?;
214        }
215        _ => {
216            // JPEG with quality control
217            let mut output_file = std::fs::File::create(&output_path).map_err(|e| {
218                ImmichError::Io(std::io::Error::other(format!(
219                    "Failed to create output file: {}",
220                    e
221                )))
222            })?;
223
224            let encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(
225                &mut output_file,
226                spec.transform.quality,
227            );
228            resized.write_with_encoder(encoder).map_err(|e| {
229                ImmichError::Io(std::io::Error::other(format!("Failed to encode JPEG: {}", e)))
230            })?;
231        }
232    }
233
234    // Apply EXIF metadata
235    apply_exif(&output_path, &spec.exif, spec.transform.strip_dimensions)?;
236
237    Ok(output_path)
238}
239
240/// Generate a test video with specified dimensions.
241fn generate_video(
242    filename: &str,
243    output_dir: &Path,
244    width: Option<u32>,
245    height: Option<u32>,
246) -> Result<PathBuf> {
247    let output_path = output_dir.join(filename);
248
249    let w = width.unwrap_or(1920);
250    let h = height.unwrap_or(1080);
251    let size = format!("{}x{}", w, h);
252
253    let output = Command::new("ffmpeg")
254        .args([
255            "-y",
256            "-f",
257            "lavfi",
258            "-i",
259            &format!("color=c=blue:s={}:d=1", size),
260            "-c:v",
261            "libx264",
262            "-pix_fmt",
263            "yuv420p",
264            output_path.to_string_lossy().as_ref(),
265        ])
266        .output()
267        .map_err(|e| {
268            ImmichError::Io(std::io::Error::other(format!(
269                "Failed to run ffmpeg: {}. Is ffmpeg installed?",
270                e
271            )))
272        })?;
273
274    if !output.status.success() {
275        let stderr = String::from_utf8_lossy(&output.stderr);
276        return Err(ImmichError::Io(std::io::Error::other(format!(
277            "ffmpeg failed: {}",
278            stderr
279        ))));
280    }
281
282    Ok(output_path)
283}
284
285/// Apply EXIF metadata to an image using exiftool CLI.
286fn apply_exif(path: &Path, exif: &ExifSpec, strip_dimensions: bool) -> Result<()> {
287    let mut args: Vec<String> = vec!["-overwrite_original".to_string()];
288
289    // GPS coordinates
290    if let Some((lat, lon)) = exif.gps {
291        let lat_ref = if lat >= 0.0 { "N" } else { "S" };
292        let lon_ref = if lon >= 0.0 { "E" } else { "W" };
293        args.push(format!("-GPSLatitude={}", lat.abs()));
294        args.push(format!("-GPSLatitudeRef={}", lat_ref));
295        args.push(format!("-GPSLongitude={}", lon.abs()));
296        args.push(format!("-GPSLongitudeRef={}", lon_ref));
297    }
298
299    // Datetime
300    if let Some(dt) = &exif.datetime {
301        let formatted = dt.format("%Y:%m:%d %H:%M:%S").to_string();
302        args.push(format!("-DateTimeOriginal={}", formatted));
303    }
304
305    // Timezone
306    if let Some(tz) = &exif.timezone {
307        args.push(format!("-OffsetTimeOriginal={}", tz));
308    }
309
310    // Camera info
311    if let Some(make) = &exif.camera_make {
312        args.push(format!("-Make={}", make));
313    }
314    if let Some(model) = &exif.camera_model {
315        args.push(format!("-Model={}", model));
316    }
317
318    // Description
319    if let Some(desc) = &exif.description {
320        args.push(format!("-ImageDescription={}", desc));
321    }
322
323    // Strip dimension EXIF if requested
324    if strip_dimensions {
325        args.push("-ImageWidth=".to_string());
326        args.push("-ExifImageWidth=".to_string());
327        args.push("-ImageHeight=".to_string());
328        args.push("-ExifImageHeight=".to_string());
329    }
330
331    // Only run exiftool if we have args to apply
332    if args.len() > 1 {
333        args.push(path.to_string_lossy().to_string());
334
335        let output = Command::new("exiftool")
336            .args(&args)
337            .output()
338            .map_err(|e| {
339                ImmichError::Io(std::io::Error::other(format!(
340                    "Failed to run exiftool: {}. Is exiftool installed?",
341                    e
342                )))
343            })?;
344
345        if !output.status.success() {
346            let stderr = String::from_utf8_lossy(&output.stderr);
347            return Err(ImmichError::Io(std::io::Error::other(format!(
348                "exiftool failed: {}",
349                stderr
350            ))));
351        }
352    }
353
354    Ok(())
355}
356
357#[cfg(test)]
358mod tests {
359    use super::*;
360
361    #[test]
362    fn test_transform_spec_builder() {
363        let spec = TransformSpec::new("base_landscape.jpg")
364            .with_size(1000, 750)
365            .with_quality(90);
366
367        assert_eq!(spec.base_image, "base_landscape.jpg");
368        assert_eq!(spec.width, Some(1000));
369        assert_eq!(spec.height, Some(750));
370        assert_eq!(spec.quality, 90);
371    }
372
373    #[test]
374    fn test_transform_spec_scale() {
375        let spec = TransformSpec::new("base_portrait.jpg").with_scale(50);
376
377        assert_eq!(spec.width, Some(50));
378        assert_eq!(spec.height, None);
379    }
380}