Add support for S3-backed image storage
This commit is contained in:
1 parent
7fcb2eb747
commit
8763a8e74a
55 files changed
+2611
-111
No files matched your search
@@ -0,0 +1,131 @@
|
||||
use std::{
|
||||
collections::BTreeSet,
|
||||
fs,
|
||||
io::Cursor,
|
||||
path::{Path, PathBuf},
|
||||
};
|
||||
|
||||
use image::{ImageFormat, RgbaImage};
|
||||
|
||||
use super::{prepare_image, ImageError, MODEL_BYTES};
|
||||
|
||||
const NSFW_IMAGES: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/nsfw");
|
||||
const NON_NSFW_IMAGES: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/fixtures/non_nsfw");
|
||||
|
||||
fn fixture_pairs(directory: &Path) -> Vec<(PathBuf, PathBuf)> {
|
||||
let mut indices = BTreeSet::new();
|
||||
for entry in fs::read_dir(directory).expect("Could not read image fixture directory") {
|
||||
let entry = entry.expect("Could not read image fixture entry");
|
||||
let path = entry.path();
|
||||
if path.extension().and_then(|extension| extension.to_str()) != Some("png") {
|
||||
continue;
|
||||
}
|
||||
let name = path.file_name().unwrap().to_str().expect("Fixture name is not valid UTF-8");
|
||||
let index = name
|
||||
.strip_prefix("image_")
|
||||
.and_then(|name| name.strip_suffix("a.png").or_else(|| name.strip_suffix("b.png")))
|
||||
.filter(|index| index.parse::<u32>().is_ok_and(|index| index > 0))
|
||||
.unwrap_or_else(|| {
|
||||
panic!("Invalid fixture name {name}; expected image_1a.png/image_1b.png")
|
||||
});
|
||||
indices.insert(index.to_string());
|
||||
}
|
||||
|
||||
assert!(
|
||||
!indices.is_empty(),
|
||||
"No paired fixtures in {}. Add images with scripts/split_nsfw_test_image.py.",
|
||||
directory.display()
|
||||
);
|
||||
indices
|
||||
.into_iter()
|
||||
.map(|index| {
|
||||
let first = directory.join(format!("image_{index}a.png"));
|
||||
let second = directory.join(format!("image_{index}b.png"));
|
||||
assert!(
|
||||
first.is_file() && second.is_file(),
|
||||
"Both fixture parts must exist: {} and {}",
|
||||
first.display(),
|
||||
second.display()
|
||||
);
|
||||
(first, second)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn recombine_parts(first: &RgbaImage, second: &RgbaImage) -> Result<RgbaImage, &'static str> {
|
||||
if first.dimensions() != second.dimensions() {
|
||||
return Err("Fixture parts must have identical dimensions");
|
||||
}
|
||||
let pixels = first.as_raw().iter().zip(second.as_raw()).map(|(a, b)| a ^ b).collect();
|
||||
RgbaImage::from_raw(first.width(), first.height(), pixels).ok_or("Invalid fixture pixel data")
|
||||
}
|
||||
|
||||
fn check_directory(directory: &str, expected_nsfw: bool) {
|
||||
let pairs = fixture_pairs(Path::new(directory));
|
||||
let model = nsfw::create_model(MODEL_BYTES).expect("Could not load bundled NSFW model");
|
||||
for (first, second) in pairs {
|
||||
let a = image::open(&first).expect("Could not decode first fixture part").to_rgba8();
|
||||
let b = image::open(&second).expect("Could not decode second fixture part").to_rgba8();
|
||||
let reconstructed = recombine_parts(&a, &b).expect("Could not recombine fixture parts");
|
||||
let mut bytes = Cursor::new(Vec::new());
|
||||
reconstructed
|
||||
.write_to(&mut bytes, ImageFormat::Png)
|
||||
.expect("Could not encode reconstructed fixture");
|
||||
|
||||
// Use the production path, including flattening transparency and the same
|
||||
// combined category threshold. Never write the reconstructed image to disk.
|
||||
let result = prepare_image(Some(&model), bytes.get_ref());
|
||||
if expected_nsfw {
|
||||
assert_eq!(
|
||||
result.as_ref().err(),
|
||||
Some(&ImageError::Nsfw),
|
||||
"Expected {} to be blocked as NSFW, got {:?}",
|
||||
first.display(),
|
||||
result.as_ref().err()
|
||||
);
|
||||
} else {
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"Expected {} to be accepted, got {:?}",
|
||||
first.display(),
|
||||
result.as_ref().err()
|
||||
);
|
||||
}
|
||||
// The toggle must allow either category, while retaining image validation.
|
||||
assert!(
|
||||
prepare_image(None, bytes.get_ref()).is_ok(),
|
||||
"Expected {} to be accepted with moderation disabled",
|
||||
first.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "Add local split fixtures to tests/fixtures/nsfw, then run with --ignored"]
|
||||
fn nsfw_directory_images_are_blocked() {
|
||||
check_directory(NSFW_IMAGES, true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[ignore = "Add local split fixtures to tests/fixtures/non_nsfw, then run with --ignored"]
|
||||
fn non_nsfw_directory_images_are_allowed() {
|
||||
check_directory(NON_NSFW_IMAGES, false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recombining_parts_restores_every_pixel_and_alpha_channel() {
|
||||
let original = [0, 17, 254, 255, 45, 99, 128, 0];
|
||||
let random = vec![18, 201, 50, 34, 166, 76, 11, 220];
|
||||
let masked = original.iter().zip(&random).map(|(pixel, mask)| pixel ^ mask).collect();
|
||||
let first = RgbaImage::from_raw(2, 1, random).unwrap();
|
||||
let second = RgbaImage::from_raw(2, 1, masked).unwrap();
|
||||
assert_eq!(recombine_parts(&first, &second).unwrap().as_raw(), &original);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_parts_with_different_dimensions() {
|
||||
assert_eq!(
|
||||
recombine_parts(&RgbaImage::new(1, 1), &RgbaImage::new(2, 1)),
|
||||
Err("Fixture parts must have identical dimensions")
|
||||
);
|
||||
}
|
||||
Reference in new issue
Block a user