From 7ea0bb0dd281536c5f20a0ca97c040f8aa95ff5e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 15 May 2026 10:49:56 -0400 Subject: [PATCH] gpt-image-2 with HD options; refactor image omni gen and pricing pipelines (#1523) --- .../src/configs/omni_gen/image_models.rs | 9 + .../cost/image/omni_gen_image_cost_handler.rs | 101 +- .../generate/image/distill_helper/mod.rs | 3 - .../image/distill_helper/request_to_costs.rs | 25 - .../image/distill_helper/request_to_plan.rs | 20 - .../generate/image/distill_image_request.rs | 139 -- .../hydrate_to_router_request.rs | 82 +- .../endpoints/omni_gen/generate/image/mod.rs | 10 +- .../image/omni_gen_image_generate_handler.rs | 106 +- .../generate/image/pipeline_result.rs | 8 + .../generate/image/pipeline_v1/mod.rs | 1 + .../image/pipeline_v1/run_pipeline_v1.rs | 121 ++ .../generate/image/pipeline_v2/mod.rs | 1 + .../image/pipeline_v2/run_pipeline_v2.rs | 187 +++ .../generate/image/resolve_media_tokens.rs | 74 - .../omni_gen/generate/image/tests/mod.rs | 13 - .../generate/image/tests/test_flux_1_dev.rs | 421 ------ .../image/tests/test_flux_1_schnell.rs | 529 ------- .../generate/image/tests/test_flux_pro_1p1.rs | 412 ----- .../image/tests/test_flux_pro_1p1_ultra.rs | 417 ------ .../generate/image/tests/test_gpt_image_1.rs | 823 ---------- .../image/tests/test_gpt_image_1p5.rs | 721 --------- .../generate/image/tests/test_gpt_image_2.rs | 655 -------- .../generate/image/tests/test_nano_banana.rs | 948 ------------ .../image/tests/test_nano_banana_2.rs | 1321 ----------------- .../image/tests/test_nano_banana_pro.rs | 1211 --------------- .../generate/image/tests/test_seedream_4.rs | 613 -------- .../generate/image/tests/test_seedream_4p5.rs | 610 -------- .../image/tests/test_seedream_5_lite.rs | 610 -------- .../util/lookup/lookup_image_urls_as_map.rs | 1 + .../lookup_image_urls_as_optional_list.rs | 1 + .../lookup/lookup_media_file_urls_as_map.rs | 1 + ...kup_media_files_as_cdn_url_list_and_map.rs | 96 ++ .../storyteller_web/src/util/lookup/mod.rs | 3 +- frontend/apps/artcraft-webapp/src/main.tsx | 2 +- frontend/apps/artcraft-website/src/main.tsx | 2 +- 36 files changed, 603 insertions(+), 9694 deletions(-) delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/mod.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_costs.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_plan.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_image_request.rs rename crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/{distill_helper => }/hydrate_to_router_request.rs (100%) create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/mod.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/run_pipeline_v1.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/mod.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/resolve_media_tokens.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/mod.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_dev.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_schnell.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1_ultra.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1p5.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_2.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_2.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_pro.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4p5.rs delete mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_5_lite.rs create mode 100644 crates/service/web/storyteller_web/src/util/lookup/lookup_media_files_as_cdn_url_list_and_map.rs diff --git a/crates/service/web/storyteller_web/src/configs/omni_gen/image_models.rs b/crates/service/web/storyteller_web/src/configs/omni_gen/image_models.rs index 35ae468938..9c9e2fae86 100644 --- a/crates/service/web/storyteller_web/src/configs/omni_gen/image_models.rs +++ b/crates/service/web/storyteller_web/src/configs/omni_gen/image_models.rs @@ -270,6 +270,15 @@ fn build_omni_gen_image_models() -> Vec { CommonQuality::Low, ]), default_quality: Some(CommonQuality::High), + // NOTE: Gpt-Image-2 does not have "resolution" natively. We're emulating this. + resolution_options: Some(vec![ + CommonResolution::OneK, + CommonResolution::TwoK, + CommonResolution::ThreeK, + CommonResolution::FourK, + ]), + // NOTE: Gpt-Image-2 does not have "resolution" natively. We're emulating this. + resolution_default: Some(CommonResolution::OneK), batch_size_min: Some(1), batch_size_max: Some(4), batch_size_default: Some(1), diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/cost/image/omni_gen_image_cost_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/cost/image/omni_gen_image_cost_handler.rs index 81ea5691d6..31f9c79f40 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/cost/image/omni_gen_image_cost_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/cost/image/omni_gen_image_cost_handler.rs @@ -4,11 +4,14 @@ use actix_web::web::Json; use actix_web::{web, HttpRequest}; use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; use artcraft_api_defs::omni_gen::cost_response::omni_gen_image_cost_response::OmniGenImageCostResponse; +use artcraft_router::api::image_list_ref::ImageListRef; use artcraft_router::api::provider::Provider; +use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; +use artcraft_router::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate; use log::warn; use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; -use crate::http_server::endpoints::omni_gen::generate::image::distill_helper::hydrate_to_router_request::hydrate_to_router_request; +use crate::http_server::endpoints::omni_gen::generate::image::hydrate_to_router_request::hydrate_to_router_request; use crate::state::server_state::ServerState; /// Estimate the cost of an image generation. @@ -29,15 +32,14 @@ pub async fn omni_gen_image_cost_handler( _server_state: web::Data>, ) -> Result, AdvancedCommonWebError> { let mut generate_request = hydrate_to_router_request(&request)?; - generate_request.provider = Provider::Artcraft; + generate_request.provider = Provider::Artcraft; // NB: Explicitly spell this out. - let plan = generate_request.build() - .map_err(|e| { - warn!("Failed to build cost plan: {}", e); - AdvancedCommonWebError::from_error(e) - })?; - - let estimate = plan.estimate_costs(); + // TODO(bt,2026-05-15): Get rid of this as soon as we can (there are even worse hacks below). + let estimate = if should_use_pipeline_v2(&generate_request) { + estimate_pipeline_v2_cost(&generate_request)? + } else { + estimate_pipeline_v1_cost(&generate_request)? + }; Ok(Json(OmniGenImageCostResponse { success: true, @@ -49,3 +51,84 @@ pub async fn omni_gen_image_cost_handler( has_watermark: estimate.has_watermark, })) } + +fn should_use_pipeline_v2( + generate_request: &GenerateImageRequestBuilder, +) -> bool { + // Early migrate anything that uses FAL v2: + let mut builder = generate_request.clone(); + builder.provider = Provider::Fal; + + if builder.use_new_builder() { + return true; + } + + // Fallback + generate_request.use_new_builder() +} + +fn estimate_pipeline_v2_cost( + generate_request: &GenerateImageRequestBuilder, +) -> Result { + let mut builder = generate_request.clone(); + + builder.provider = Provider::Artcraft; + + match builder.clone().build2() { + Ok(result) => match result.estimate_cost() { + Ok(cost) => return Ok(cost), + Err(err) => { + warn!("Failed to estimate image cost for pipeline v2 (artcraft): {}", err); + } + }, + Err(err) => { + warn!("Failed to build image cost request for pipeline v2 (artcraft): {}", err); + } + } + + // TODO(bt,2026-05-15): This is a horrible hack for models that aren't fully migrated + // to pipeline_v2 (don't have artcraft cost calculators). This is an awful hack we need + // to kill in the future. + + builder.provider = Provider::Fal; + + // Fal cost estimation doesn't work with media file tokens. + // We don't need to look them up (though, technically, image size does impact cost with + // some providers). Let's just fill the list with dummy URLs in that case. + if let Some(ImageListRef::MediaFileTokens(tokens)) = builder.image_inputs.as_ref() { + builder.image_inputs = Some(ImageListRef::Urls(vec![ + "https://example.com/image.png".to_string(); + tokens.len() + ])); + } + + let mut cost = builder.build2() + .map_err(|e| { + warn!("Failed to build image cost request for pipeline v2 (fal): {}", e); + AdvancedCommonWebError::from_error(e) + })? + .estimate_cost() + .map_err(|e| { + warn!("Failed to estimate image cost for pipeline v2 (fal): {}", e); + AdvancedCommonWebError::from_error(e) + })?; + + cost.cost_in_credits = cost.cost_in_usd_cents; + + Ok(cost) +} + +fn estimate_pipeline_v1_cost( + generate_request: &GenerateImageRequestBuilder, +) -> Result { + let mut builder = generate_request.clone(); + builder.provider = Provider::Artcraft; + + let plan = builder.build() + .map_err(|e| { + warn!("Failed to build image cost plan for pipeline v1: {}", e); + AdvancedCommonWebError::from_error(e) + })?; + + Ok(plan.estimate_costs()) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/mod.rs deleted file mode 100644 index 706ab177f0..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub(crate) mod hydrate_to_router_request; -pub(crate) mod request_to_costs; -pub(crate) mod request_to_plan; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_costs.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_costs.rs deleted file mode 100644 index 1e84fe2aa3..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_costs.rs +++ /dev/null @@ -1,25 +0,0 @@ -use artcraft_router::api::provider::Provider; -use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; -use artcraft_router::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate; -use log::warn; - -use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; - -/// Estimate costs for an image generation request. -/// Always uses the Artcraft provider for costing regardless of the execution provider. -/// (Because we're the ones doing the billing.) -pub fn request_to_costs( - request: &GenerateImageRequestBuilder, -) -> Result { - let cost_request = GenerateImageRequestBuilder { - provider: Provider::Artcraft, - ..request.clone() - }; - - let plan = cost_request.build().map_err(|e| { - warn!("Failed to build cost plan: {}", e); - AdvancedCommonWebError::from_error(e) - })?; - - Ok(plan.estimate_costs()) -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_plan.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_plan.rs deleted file mode 100644 index 90b34132ca..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_plan.rs +++ /dev/null @@ -1,20 +0,0 @@ -use artcraft_router::api::provider::Provider; -use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; -use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; -use log::warn; - -use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; - -/// Build an image generation plan from a transformed request. -/// Overrides the provider to Fal for execution (costing uses Artcraft separately). -pub fn request_to_plan( - request: &mut GenerateImageRequestBuilder, -) -> Result { - request.provider = Provider::Fal; - - request.build() - .map_err(|e| { - warn!("Failed to build image generation plan: {}", e); - AdvancedCommonWebError::from_error(e) - }) -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_image_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_image_request.rs deleted file mode 100644 index e0b1b33177..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_image_request.rs +++ /dev/null @@ -1,139 +0,0 @@ -//! Distillation step for the omni-gen image endpoints. -//! -//! Takes a raw [`OmniGenImageCostAndGenerateRequest`] plus a (pre-computed) -//! `MediaFileToken -> Url` map and produces a fully self-contained -//! [`DistilledImageRequest`] holding: -//! - the (private) router request, kept for inspection in tests / debugging -//! - the [`ImageGenerationCostEstimate`] (Artcraft provider, what we bill on) -//! - the [`ImageGenerationPlan`] (Fal provider, what we actually execute) - -use std::collections::HashMap; - -use log::warn; -use url::Url; - -use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; -use artcraft_router::api::image_list_ref::ImageListRef; -use artcraft_router::api::provider::Provider; -use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; -use artcraft_router::generate::generate_image::image_generation_cost_estimate::ImageGenerationCostEstimate; -use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; -use tokens::tokens::media_files::MediaFileToken; - -use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; - -use super::distill_helper::hydrate_to_router_request::hydrate_to_router_request; - -/// Self-contained, owned representation of a fully-distilled omni-gen image -/// request: the router request, the bill-on cost estimate, and the executable -/// plan, all in one place. -pub struct DistilledImageRequest { - /// The fully-built router request. - pub request: GenerateImageRequestBuilder, - - /// Cost estimate as computed by the Artcraft provider — this is what we bill on. - pub cost: ImageGenerationCostEstimate, - - /// Execution plan as computed by the Fal provider — what we hand to the router. - pub plan: ImageGenerationPlan, -} - -impl DistilledImageRequest { - /// Borrow the execution plan. - pub fn plan(&self) -> &ImageGenerationPlan { - &self.plan - } - - /// Borrow the underlying router request. Useful for tests / debugging. - #[allow(dead_code)] - pub(crate) fn request(&self) -> &GenerateImageRequestBuilder { - &self.request - } -} - -/// Build a [`DistilledImageRequest`] from a raw API request and a pre-computed -/// `MediaFileToken -> Url` hydration map. -/// -/// The hydration map should already have been resolved by the caller (via -/// `lookup_image_urls_as_map` or equivalent) — distillation does no I/O. -/// -/// Returns `BadInputWithSimpleMessage` if any image media token referenced by -/// the request is missing from the hydration map. -pub fn distill_image_request( - request: &OmniGenImageCostAndGenerateRequest, - media_file_hydration_map: Option<&HashMap>, -) -> Result { - // 1. Convert the raw API request into a router request. - let mut initial = hydrate_to_router_request(request)?; - - // 2. If we have media tokens, resolve them to URLs for the execution request. - let image_input_urls = build_image_input_urls( - request.image_media_tokens.as_ref(), - media_file_hydration_map, - )?; - - // 3. Build the execution request with resolved URLs and Fal provider. - if let Some(urls) = image_input_urls { - initial.image_inputs = Some(ImageListRef::Urls(urls)); - } - initial.provider = Provider::Fal; - let fal_request = initial; - - // 4. Cost estimate. Always use the Artcraft provider for billing regardless - // of which provider executes the request. - let cost_request = GenerateImageRequestBuilder { - provider: Provider::Artcraft, - ..fal_request.clone() - }; - let cost_plan = cost_request.build().map_err(|e| { - warn!("Failed to build cost plan during distillation: {}", e); - AdvancedCommonWebError::from_error(e) - })?; - let cost: ImageGenerationCostEstimate = cost_plan.estimate_costs(); - - // 5. Execution plan. - let plan = fal_request.build().map_err(|e| { - warn!("Failed to build image generation plan during distillation: {}", e); - AdvancedCommonWebError::from_error(e) - })?; - - Ok(DistilledImageRequest { - request: fal_request, - cost, - plan, - }) -} - -/// Build the owned `Vec` of image input URLs from the raw token list + -/// hydration map. Errors out if any token referenced by the request is missing -/// from the hydration map. -fn build_image_input_urls( - image_media_tokens: Option<&Vec>, - hydration_map: Option<&HashMap>, -) -> Result>, AdvancedCommonWebError> { - let tokens = match image_media_tokens { - Some(tokens) if !tokens.is_empty() => tokens, - _ => return Ok(None), - }; - - let map = hydration_map.ok_or_else(|| { - AdvancedCommonWebError::BadInputWithSimpleMessage( - "image_media_tokens supplied but no hydration map was provided".to_string(), - ) - })?; - - let mut urls: Vec = Vec::with_capacity(tokens.len()); - for token in tokens { - match map.get(token) { - Some(url) => urls.push(url.to_string()), - None => { - return Err(AdvancedCommonWebError::BadInputWithSimpleMessage(format!( - "Image media token not found in hydration map: {:?}", - token - ))); - } - } - } - - Ok(Some(urls)) -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/hydrate_to_router_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/hydrate_to_router_request.rs similarity index 100% rename from crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/hydrate_to_router_request.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/hydrate_to_router_request.rs index 6e2a0b474e..76006ddad1 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/hydrate_to_router_request.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/hydrate_to_router_request.rs @@ -14,6 +14,47 @@ use enums::common::generation::common_resolution::CommonResolution as CommonReso use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; +pub fn hydrate_to_router_request( + request: &OmniGenImageCostAndGenerateRequest, +) -> Result { + let api_model = request.model.as_ref() + .ok_or_else(|| AdvancedCommonWebError::BadInputWithSimpleMessage( + "model is required".to_string(), + ))?; + + let model = convert_model(api_model)?; + + let aspect_ratio = request.aspect_ratio.as_ref() + .map(convert_aspect_ratio) + .transpose()?; + + let resolution = request.resolution.as_ref() + .map(convert_resolution) + .transpose()?; + + let quality = request.quality.as_ref() + .map(convert_quality) + .transpose()?; + + Ok(GenerateImageRequestBuilder { + model, + provider: Provider::Artcraft, + prompt: request.prompt.clone(), + image_inputs: request.image_media_tokens.clone() + .map(ImageListRef::MediaFileTokens), + resolution, + aspect_ratio, + quality, + image_batch_count: request.image_batch_count, + horizontal_angle: request.adjust_horizontal_angle, + vertical_angle: request.adjust_vertical_angle, + zoom: request.adjust_zoom, + request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade, + generation_mode_mismatch_strategy: None, + idempotency_token: request.idempotency_token.clone(), + }) +} + fn convert_model( model: &CommonImageModelEnum, ) -> Result { @@ -57,44 +98,3 @@ fn convert_quality( ) }) } - -pub fn hydrate_to_router_request( - request: &OmniGenImageCostAndGenerateRequest, -) -> Result { - let api_model = request.model.as_ref() - .ok_or_else(|| AdvancedCommonWebError::BadInputWithSimpleMessage( - "model is required".to_string(), - ))?; - - let model = convert_model(api_model)?; - - let aspect_ratio = request.aspect_ratio.as_ref() - .map(convert_aspect_ratio) - .transpose()?; - - let resolution = request.resolution.as_ref() - .map(convert_resolution) - .transpose()?; - - let quality = request.quality.as_ref() - .map(convert_quality) - .transpose()?; - - Ok(GenerateImageRequestBuilder { - model, - provider: Provider::Artcraft, - prompt: request.prompt.clone(), - image_inputs: request.image_media_tokens.clone() - .map(ImageListRef::MediaFileTokens), - resolution, - aspect_ratio, - quality, - image_batch_count: request.image_batch_count, - horizontal_angle: request.adjust_horizontal_angle, - vertical_angle: request.adjust_vertical_angle, - zoom: request.adjust_zoom, - request_mismatch_mitigation_strategy: RequestMismatchMitigationStrategy::PayMoreUpgrade, - generation_mode_mismatch_strategy: None, - idempotency_token: request.idempotency_token.clone(), - }) -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs index 8fdd478417..653af91719 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs @@ -1,7 +1,5 @@ pub mod omni_gen_image_generate_handler; -pub(crate) mod distill_helper; -pub(crate) mod distill_image_request; -mod resolve_media_tokens; - -#[cfg(test)] -mod tests; +pub(crate) mod hydrate_to_router_request; +pub(crate) mod pipeline_result; +pub mod pipeline_v1; +pub mod pipeline_v2; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs index 513d546958..6e9025695b 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs @@ -1,16 +1,12 @@ -use std::collections::HashMap; use std::sync::Arc; use actix_web::web::Json; use actix_web::{web, HttpRequest}; use log::{error, info, warn}; use sqlx::Acquire; -use url::Url; use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; use artcraft_api_defs::omni_gen::generate_response::omni_gen_image_generate_response::OmniGenImageGenerateResponse; -use artcraft_router::client::router_client::RouterClient; -use artcraft_router::client::router_fal_client::RouterFalClient; use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse; use enums::by_table::debug_logs::debug_log_type::DebugLogType; use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType; @@ -31,18 +27,16 @@ use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_it insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem, }; use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs}; -use tokens::tokens::generic_inference_jobs::InferenceJobToken; -use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::non_unique::debug_logs_event_token::DebugLogEventToken; -use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error; use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test; +use crate::http_server::endpoints::omni_gen::generate::image::hydrate_to_router_request::hydrate_to_router_request; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::run_pipeline_v1::{run_pipeline_v1, RunPipelineV1Args}; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v2::run_pipeline_v2::{run_pipeline_v2, should_use_pipeline_v2, RunPipelineV2Args}; use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; use crate::state::server_state::ServerState; -use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map; - -use super::distill_image_request::distill_image_request; +use crate::util::lookup::lookup_media_files_as_cdn_url_list_and_map::lookup_media_files_as_cdn_url_list_and_map; /// Generate an image using the omni-gen unified endpoint. #[utoipa::path( @@ -112,52 +106,18 @@ pub async fn omni_gen_image_generate_handler( })?; // ==================== RESOLVE MEDIA TOKENS ==================== // - // Look up media file tokens BEFORE distilling. distill_image_request takes - // a pre-computed `MediaFileToken -> Url` map and does no I/O of its own. - let media_file_hydration_map: Option> = match request.image_media_tokens.as_ref() { - Some(tokens) if !tokens.is_empty() => { - info!("Resolving {} image media tokens to CDN URLs", tokens.len()); - let raw = lookup_image_urls_as_map( - &http_request, - &mut mysql_connection, - server_state.server_environment, - tokens, - ).await?; - let parsed: HashMap = raw.into_iter() - .filter_map(|(token, url_str)| match Url::parse(&url_str) { - Ok(url) => Some((token, url)), - Err(err) => { - warn!("Failed to parse media file URL {:?}: {:?}", url_str, err); - None - } - }) - .collect(); - Some(parsed) - } - _ => None, - }; + // Look up media file tokens BEFORE distilling. Pipeline execution should not do I/O. + let resolved_media = lookup_media_files_as_cdn_url_list_and_map( + &http_request, + &mut mysql_connection, + server_state.server_environment, + request.image_media_tokens.as_deref().unwrap_or(&[]), + ).await?; - // ==================== DISTILL ==================== // + // ==================== HYDRATE ROUTER REQUEST ==================== // - let distilled = distill_image_request(&request, media_file_hydration_map.as_ref())?; - - // ==================== BILLING ==================== // - - let cost = distilled.cost.cost_in_credits.unwrap_or(0); - - info!("Charging wallet: {} credits", cost); - - let apriori_job_token = InferenceJobToken::generate(); - - if cost > 0 { - attempt_wallet_deduction_else_common_web_error( - user_token, - Some(apriori_job_token.as_str()), - cost, - &mut mysql_connection, - ).await?; - } + let router_builder = hydrate_to_router_request(&request)?; // ==================== DEBUG LOG: HTTP REQUEST ==================== // @@ -172,25 +132,31 @@ pub async fn omni_gen_image_generate_handler( warn!("Failed to insert HTTP request debug log: {:?}", err); } - // ==================== EXECUTE GENERATION ==================== // + // ==================== PIPELINE DISPATCH ==================== // - let fal_client = RouterFalClient::new( - server_state.fal.api_key.clone(), - server_state.fal.webhook_url.clone(), - ); - - let router_client = RouterClient::Fal(fal_client); - - let generation_response = distilled.plan().generate_image(&router_client) - .await - .map_err(|e| { - warn!("Image generation failed: {:?}", e); - AdvancedCommonWebError::from_error(e) - })?; + let pipeline_result = if should_use_pipeline_v2(&router_builder) { + info!("Using image pipeline v2"); + run_pipeline_v2(RunPipelineV2Args { + router_builder: &router_builder, + server_state: &server_state, + mysql_connection: &mut mysql_connection, + user_token, + resolved_media: &resolved_media, + }).await? + } else { + info!("Using image pipeline v1"); + run_pipeline_v1(RunPipelineV1Args { + router_builder: &router_builder, + server_state: &server_state, + mysql_connection: &mut mysql_connection, + user_token, + resolved_media: &resolved_media, + }).await? + }; // ==================== DEBUG LOG: FAL REQUEST ==================== // - if let GenerateImageResponse::Fal(ref fal_payload) = generation_response { + if let GenerateImageResponse::Fal(ref fal_payload) = pipeline_result.response { if let Some(ref outbound_request) = fal_payload.maybe_outbound_request { if let Err(err) = insert_debug_log(InsertDebugLogArgs { apriori_debug_log_event_token: Some(&debug_log_event_token), @@ -205,7 +171,7 @@ pub async fn omni_gen_image_generate_handler( } } - let external_job_id = match &generation_response { + let external_job_id = match &pipeline_result.response { GenerateImageResponse::Artcraft(p) => { p.inference_job_token.as_str().to_string() } @@ -291,7 +257,7 @@ pub async fn omni_gen_image_generate_handler( let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token( InsertGenericInferenceForFalWithAprioriJobTokenArgs { - apriori_job_token: &apriori_job_token, + apriori_job_token: &pipeline_result.apriori_job_token, uuid_idempotency_token: &idempotency_token, maybe_external_third_party_id: &external_job_id, fal_category: FalCategory::ImageGeneration, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs new file mode 100644 index 0000000000..80a094b166 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs @@ -0,0 +1,8 @@ +use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; + +/// Both image pipelines produce this, then the shared handler suffix writes DB rows. +pub struct ImagePipelineResult { + pub apriori_job_token: InferenceJobToken, + pub response: GenerateImageResponse, +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/mod.rs new file mode 100644 index 0000000000..06848147b6 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/mod.rs @@ -0,0 +1 @@ +pub mod run_pipeline_v1; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/run_pipeline_v1.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/run_pipeline_v1.rs new file mode 100644 index 0000000000..7944616714 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/run_pipeline_v1.rs @@ -0,0 +1,121 @@ +use log::{info, warn}; + +use artcraft_router::api::image_list_ref::ImageListRef; +use artcraft_router::api::provider::Provider; +use artcraft_router::client::router_client::RouterClient; +use artcraft_router::client::router_fal_client::RouterFalClient; +use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; +use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::users::UserToken; + +use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error; +use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_result::ImagePipelineResult; +use crate::state::server_state::ServerState; +use crate::util::lookup::lookup_media_files_as_cdn_url_list_and_map::MediaFilesAsCdnUrlListAndMap; + +pub struct RunPipelineV1Args<'a> { + pub router_builder: &'a GenerateImageRequestBuilder, + pub server_state: &'a ServerState, + pub mysql_connection: &'a mut sqlx::pool::PoolConnection, + pub user_token: &'a UserToken, + pub resolved_media: &'a MediaFilesAsCdnUrlListAndMap, +} + +pub async fn run_pipeline_v1( + args: RunPipelineV1Args<'_>, +) -> Result { + let RunPipelineV1Args { + router_builder, + server_state, + mysql_connection, + user_token, + resolved_media, + } = args; + + let hydrated_builder = apply_hydrated_media_inputs( + router_builder, + resolved_media, + ); + + let cost = estimate_cost_in_credits(&router_builder)?; + + let execution_plan = build_execution_plan(&hydrated_builder)?; + + info!("Charging wallet: {} credits", cost); + + let apriori_job_token = InferenceJobToken::generate(); + + if cost > 0 { + attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost, + mysql_connection, + ).await?; + } + + let fal_client = RouterFalClient::new( + server_state.fal.api_key.clone(), + server_state.fal.webhook_url.clone(), + ); + + let router_client = RouterClient::Fal(fal_client); + + let response = execution_plan.generate_image(&router_client) + .await + .map_err(|e| { + warn!("Image generation failed: {:?}", e); + AdvancedCommonWebError::from_error(e) + })?; + + Ok(ImagePipelineResult { + apriori_job_token, + response, + }) +} + +fn apply_hydrated_media_inputs( + router_builder: &GenerateImageRequestBuilder, + resolved_media: &MediaFilesAsCdnUrlListAndMap, +) -> GenerateImageRequestBuilder { + let mut hydrated_builder = router_builder.clone(); + + match hydrated_builder.image_inputs.as_ref() { + Some(ImageListRef::MediaFileTokens(tokens)) if !tokens.is_empty() => { + hydrated_builder.image_inputs = Some(ImageListRef::Urls( + resolved_media.ordered_url_list.clone(), + )); + }, + _ => {}, + } + + hydrated_builder +} + +fn estimate_cost_in_credits( + router_builder: &GenerateImageRequestBuilder, +) -> Result { + let mut cost_builder = router_builder.clone(); + cost_builder.provider = Provider::Artcraft; + + let cost_plan = cost_builder.build().map_err(|e| { + warn!("Failed to build image cost plan for v1 pipeline: {}", e); + AdvancedCommonWebError::from_error(e) + })?; + + Ok(cost_plan.estimate_costs().cost_in_credits.unwrap_or(0)) +} + +fn build_execution_plan( + router_builder: &GenerateImageRequestBuilder, +) -> Result { + let mut execution_builder = router_builder.clone(); + execution_builder.provider = Provider::Fal; + + execution_builder.build().map_err(|e| { + warn!("Failed to build image generation plan for v1 pipeline: {}", e); + AdvancedCommonWebError::from_error(e) + }) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/mod.rs new file mode 100644 index 0000000000..f207b21293 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/mod.rs @@ -0,0 +1 @@ +pub mod run_pipeline_v2; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs new file mode 100644 index 0000000000..17dadd5b83 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs @@ -0,0 +1,187 @@ +use log::{info, warn}; + +use artcraft_router::api::image_list_ref::ImageListRef; +use artcraft_router::api::provider::Provider; +use artcraft_router::client::router_client::RouterClient; +use artcraft_router::client::router_fal_webhook_optional_client::RouterFalWebhookOptionalClient; +use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; +use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse; +use artcraft_router::generate::generate_image_v2::image_generation_draft_context::ImageGenerationDraftContext; +use artcraft_router::generate::generate_image_v2::image_generation_draft_or_request::ImageGenerationDraftOrRequest; +use artcraft_router::generate::generate_image_v2::image_generation_request::ImageGenerationRequest; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::users::UserToken; + +use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error; +use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_result::ImagePipelineResult; +use crate::state::server_state::ServerState; +use crate::util::lookup::lookup_media_files_as_cdn_url_list_and_map::MediaFilesAsCdnUrlListAndMap; + +pub struct RunPipelineV2Args<'a> { + pub router_builder: &'a GenerateImageRequestBuilder, + pub server_state: &'a ServerState, + pub mysql_connection: &'a mut sqlx::pool::PoolConnection, + pub user_token: &'a UserToken, + pub resolved_media: &'a MediaFilesAsCdnUrlListAndMap, +} + +pub fn should_use_pipeline_v2(router_builder: &GenerateImageRequestBuilder) -> bool { + let mut execution_builder = router_builder.clone(); + execution_builder.provider = Provider::Fal; + execution_builder.use_new_builder() +} + +pub async fn run_pipeline_v2( + args: RunPipelineV2Args<'_>, +) -> Result { + let RunPipelineV2Args { + router_builder, + server_state, + mysql_connection, + user_token, + resolved_media, + } = args; + + let hydrated_builder = apply_hydrated_media_inputs( + router_builder, + resolved_media, + ); + + let cost = estimate_cost_in_credits(&router_builder)?; + + let draft_or_request = build_execution_request(&hydrated_builder)?; + + info!("Charging wallet: {} credits", cost); + + let apriori_job_token = InferenceJobToken::generate(); + + if cost > 0 { + attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost, + mysql_connection, + ).await?; + } + + let response = finalize_and_generate(draft_or_request, server_state).await?; + + Ok(ImagePipelineResult { + apriori_job_token, + response, + }) +} + +fn build_execution_request( + router_builder: &GenerateImageRequestBuilder, +) -> Result { + let mut execution_builder = router_builder.clone(); + execution_builder.provider = Provider::Fal; + + execution_builder.build2().map_err(|e| { + warn!("Failed to build2 for image v2 pipeline: {}", e); + AdvancedCommonWebError::from_error(e) + }) +} + +fn apply_hydrated_media_inputs( + router_builder: &GenerateImageRequestBuilder, + resolved_media: &MediaFilesAsCdnUrlListAndMap, +) -> GenerateImageRequestBuilder { + let mut hydrated_builder = router_builder.clone(); + + match hydrated_builder.image_inputs.as_ref() { + Some(ImageListRef::MediaFileTokens(tokens)) if !tokens.is_empty() => { + hydrated_builder.image_inputs = Some(ImageListRef::Urls( + resolved_media.ordered_url_list.clone(), + )); + }, + _ => {}, + } + + hydrated_builder +} + +fn estimate_cost_in_credits( + router_builder: &GenerateImageRequestBuilder, +) -> Result { + // TODO(bt,2026-05-15): This might not be 1:1 with new Fal costs, eg. Gpt-image-2 + let mut cost_builder = router_builder.clone(); + cost_builder.provider = Provider::Artcraft; + + // Attempt pipeline_v2 cost + match cost_builder.clone().build2() { + Err(err) => { + warn!("Failed to build image cost plan for v2 pipeline: {}", err); + }, + Ok(request) => { + let cost_plan = request.estimate_cost().map_err(|e| { + warn!("Failed to build image cost plan for v2 pipeline: {}", e); + AdvancedCommonWebError::from_error(e) + })?; + return Ok(cost_plan.cost_in_credits.unwrap_or(0)); + }, + } + + // Fall back to pipeline_v1 cost + let cost_plan = cost_builder.build().map_err(|e| { + warn!("Failed to build image cost plan for v1 pipeline: {}", e); + AdvancedCommonWebError::from_error(e) + })?; + + Ok(cost_plan.estimate_costs().cost_in_credits.unwrap_or(0)) +} + +async fn finalize_and_generate( + draft_or_request: ImageGenerationDraftOrRequest, + server_state: &ServerState, +) -> Result { + let provider = draft_or_request.get_provider(); + let client = build_router_client(provider, server_state)?; + + let request = finalize_request(draft_or_request).await?; + + request.send_request(&client) + .await + .map_err(|err| { + warn!("v2 image generation failed: {:?}", err); + AdvancedCommonWebError::from_error(err) + }) +} + +async fn finalize_request( + draft_or_request: ImageGenerationDraftOrRequest, +) -> Result { + match draft_or_request { + ImageGenerationDraftOrRequest::Request(request) => Ok(request), + ImageGenerationDraftOrRequest::Draft(draft) => { + draft.finalize(ImageGenerationDraftContext::default()) + .await + .map_err(|err| { + warn!("Failed to finalize image v2 draft: {:?}", err); + AdvancedCommonWebError::from_error(err) + }) + } + } +} + +fn build_router_client( + provider: Provider, + server_state: &ServerState, +) -> Result { + match provider { + Provider::Fal => { + let fal_client = RouterFalWebhookOptionalClient::new_with_webhook( + server_state.fal.api_key.clone(), + server_state.fal.webhook_url.clone(), + ); + Ok(RouterClient::FalWebhookOptional(fal_client)) + }, + other => { + Err(AdvancedCommonWebError::server_error_with_message( + &format!("Unsupported provider for image v2 generation: {:?}", other), + )) + }, + } +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/resolve_media_tokens.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/resolve_media_tokens.rs deleted file mode 100644 index 66433b384b..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/resolve_media_tokens.rs +++ /dev/null @@ -1,74 +0,0 @@ -use std::collections::HashMap; - -use actix_web::HttpRequest; -use log::info; -use sqlx::pool::PoolConnection; -use sqlx::MySql; - -use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; -use artcraft_router::api::image_list_ref::ImageListRef; -use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; -use tokens::tokens::media_files::MediaFileToken; - -use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; -use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map; -use server_environment::ServerEnvironment; - -/// Resolved media URLs that can be referenced by a GenerateImageRequestBuilder. -/// The owned Vecs here outlive the request so the GenerateImageRequestBuilder can borrow them. -pub struct ResolvedImageMedia { - pub url_map: HashMap, - pub image_input_urls: Vec, -} - -/// Collect all media file tokens from the raw HTTP request, query them from the database, -/// and return resolved CDN URLs. -pub async fn resolve_media_tokens( - omni_request: &OmniGenImageCostAndGenerateRequest, - http_request: &HttpRequest, - mysql_connection: &mut PoolConnection, - server_environment: ServerEnvironment, -) -> Result { - let mut all_tokens: Vec = Vec::new(); - - if let Some(tokens) = &omni_request.image_media_tokens { - all_tokens.extend(tokens.iter().cloned()); - } - - if all_tokens.is_empty() { - return Ok(ResolvedImageMedia { - url_map: HashMap::new(), - image_input_urls: Vec::new(), - }); - } - - info!("Resolving {} media file tokens to CDN URLs", all_tokens.len()); - - let url_map = lookup_image_urls_as_map( - http_request, - mysql_connection, - server_environment, - &all_tokens, - ).await?; - - let image_input_urls = omni_request.image_media_tokens.as_ref() - .map(|tokens| tokens.iter().filter_map(|t| url_map.get(t).cloned()).collect()) - .unwrap_or_default(); - - Ok(ResolvedImageMedia { - url_map, - image_input_urls, - }) -} - -/// Apply resolved media URLs to a GenerateImageRequestBuilder, replacing MediaFileToken refs with URL refs. -pub fn apply_resolved_media( - request: &mut GenerateImageRequestBuilder, - resolved: &ResolvedImageMedia, -) { - if !resolved.image_input_urls.is_empty() { - request.image_inputs = Some(ImageListRef::Urls(resolved.image_input_urls.clone())); - } else { - request.image_inputs = None; - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/mod.rs deleted file mode 100644 index e08d5087a8..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/mod.rs +++ /dev/null @@ -1,13 +0,0 @@ -mod test_flux_1_dev; -mod test_flux_1_schnell; -mod test_flux_pro_1p1; -mod test_flux_pro_1p1_ultra; -mod test_gpt_image_1; -mod test_gpt_image_1p5; -mod test_gpt_image_2; -mod test_nano_banana; -mod test_nano_banana_2; -mod test_nano_banana_pro; -mod test_seedream_4; -mod test_seedream_4p5; -mod test_seedream_5_lite; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_dev.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_dev.rs deleted file mode 100644 index 8a937aea11..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_dev.rs +++ /dev/null @@ -1,421 +0,0 @@ -//! Tests for omni-gen against the Flux 1 Dev model. -//! -//! These exercise [`distill_image_request`] end-to-end without firing any -//! provider calls — both the cost estimate (Artcraft provider, what we bill -//! on) and the execution plan (Fal provider, what we'd send to fal.ai) are -//! computed synchronously from the request. -//! -//! Flux 1 Dev is text-to-image only, so no image refs are exercised. - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - /// Build a minimal Flux 1 Dev omni request with only the fields we care about - /// in tests filled in. - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - image_batch_count: Option, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::Flux1Dev), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens: None, - resolution: None, - aspect_ratio, - quality: None, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - /// Distill a request and panic with a descriptive message on failure. - /// Flux 1 Dev never uses image refs, so we always pass an empty hydration map. - fn distill(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for Flux 1 Dev") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST - // ──────────────────────────────────────────────────────────────────────────── - // - // Pricing comes from `estimate_image_cost_artcraft_flux_1_dev`: - // $0.02 / image (2 credits per image; 1 credit == 1 USD cent) - // - // The artcraft cost path is what omni bills on, so this matches what the - // user is actually charged. - // - // (Note: the Fal cost path would also be 3¢/image for flux_1_dev — but the - // omni endpoint always uses the Artcraft path for billing.) - mod cost { - use super::*; - - fn cost_for_batch(image_batch_count: Option) -> (Option, Option) { - let request = make_request(Some("a happy puppy"), None, image_batch_count); - let distilled = distill(&request); - ( - distilled.cost.cost_in_credits, - distilled.cost.cost_in_usd_cents, - ) - } - - #[test] - fn default_batch_costs_2_cents() { - // No batch specified → defaults to 1 image → 2¢ - let (credits, cents) = cost_for_batch(None); - assert_eq!(credits, Some(2)); - assert_eq!(cents, Some(2)); - } - - #[test] - fn batch_of_one_costs_2_cents() { - let (credits, cents) = cost_for_batch(Some(1)); - assert_eq!(credits, Some(2)); - assert_eq!(cents, Some(2)); - } - - #[test] - fn batch_of_two_costs_4_cents() { - let (credits, cents) = cost_for_batch(Some(2)); - assert_eq!(credits, Some(4)); - assert_eq!(cents, Some(4)); - } - - #[test] - fn batch_of_three_costs_6_cents() { - let (credits, cents) = cost_for_batch(Some(3)); - assert_eq!(credits, Some(6)); - assert_eq!(cents, Some(6)); - } - - #[test] - fn batch_of_four_costs_8_cents() { - let (credits, cents) = cost_for_batch(Some(4)); - assert_eq!(credits, Some(8)); - assert_eq!(cents, Some(8)); - } - - #[test] - fn batch_above_max_clamps_and_costs_8_cents() { - // PayMoreUpgrade strategy clamps batches >4 down to the max of 4. - let (credits, cents) = cost_for_batch(Some(7)); - assert_eq!(credits, Some(8)); - assert_eq!(cents, Some(8)); - } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let aspect_ratios = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::SquareHd), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::TallNineBySixteen), - Some(CommonAspectRatio::Auto), - ]; - for ar in aspect_ratios { - let request = make_request(Some("p"), ar, Some(2)); - let distilled = distill(&request); - assert_eq!( - distilled.cost.cost_in_usd_cents, - Some(4), - "expected 4¢ regardless of aspect ratio (got {:?})", - ar, - ); - } - } - - #[test] - fn cost_is_independent_of_prompt() { - let with_prompt = distill(&make_request(Some("a cat"), None, Some(3))); - let without_prompt = distill(&make_request(None, None, Some(3))); - assert_eq!(with_prompt.cost.cost_in_usd_cents, Some(6)); - assert_eq!(without_prompt.cost.cost_in_usd_cents, Some(6)); - } - - #[test] - fn cost_metadata_flags_are_default() { - let request = make_request(Some("p"), None, Some(1)); - let distilled = distill(&request); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - // - // The distilled plan is the Fal-side `PlanFalFlux1Dev`: - // - prompt: passthrough - // - maybe_image_url: always None (Flux 1 Dev never gets image refs in tests) - // - aspect_ratio: mapped from CommonAspectRatio with sensible fallbacks - // - num_images: mapped from image_batch_count (default 1, clamped to 1..=4) - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_flux_1_dev::{ - FalFlux1DevNumImages, PlanFalFlux1Dev, - }; - use fal_client::requests::webhook::image::text::enqueue_flux_1_dev_text_to_image_webhook::Flux1DevAspectRatio; - - /// Distill and run an assertion against the inner FalFlux1Dev plan. - /// - /// We pass the plan by reference (rather than moving it out of `distilled`) - /// because the plan's borrows are tied to data owned by `distilled` — - /// moving it out would dangle those borrows. - fn with_fal_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill(request); - match distilled.plan() { - ImageGenerationPlan::FalFlux1Dev(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalFlux1Dev, got {:?}", other), - } - } - - fn assert_plan_for_aspect_ratio( - ar: Option, - assertion: F, - ) { - with_fal_plan(&make_request(Some("p"), ar, Some(1)), assertion); - } - - fn assert_plan_for_batch( - batch: Option, - assertion: F, - ) { - with_fal_plan(&make_request(Some("p"), None, batch), assertion); - } - - // ── Aspect ratio mapping ──────────────────────────────────────────────── - - #[test] - fn default_aspect_ratio_is_square() { - assert_plan_for_aspect_ratio(None, |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::Square)); - }); - } - - #[test] - fn square_input_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::Square)); - }); - } - - #[test] - fn square_hd_input_yields_square_hd() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::SquareHd)); - }); - } - - #[test] - fn wide_4x3_yields_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn wide_16x9_yields_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn wide_alias_yields_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn tall_3x4_yields_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::PortraitThreeByFour)); - }); - } - - #[test] - fn tall_9x16_yields_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::PortraitNineBySixteen)); - }); - } - - #[test] - fn tall_alias_yields_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::PortraitNineBySixteen)); - }); - } - - // ── Aspect ratio fallback mapping ─────────────────────────────────────── - // - // Aspect ratios that Flux 1 Dev doesn't natively support get mapped to the - // closest supported ratio under the PayMoreUpgrade mismatch strategy that - // omni-gen uses. - - #[test] - fn wide_5x4_falls_back_to_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn wide_3x2_falls_back_to_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideThreeByTwo), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn wide_21x9_falls_back_to_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideTwentyOneByNine), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn tall_4x5_falls_back_to_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallFourByFive), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::PortraitThreeByFour)); - }); - } - - #[test] - fn tall_2x3_falls_back_to_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::PortraitThreeByFour)); - }); - } - - #[test] - fn tall_9x21_falls_back_to_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallNineByTwentyOne), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::PortraitNineBySixteen)); - }); - } - - // ── Auto* aspect ratios ───────────────────────────────────────────────── - - #[test] - fn auto_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::Square)); - }); - } - - #[test] - fn auto_2k_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto2k), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::Square)); - }); - } - - #[test] - fn auto_4k_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto4k), |plan| { - assert!(matches!(plan.aspect_ratio, Flux1DevAspectRatio::Square)); - }); - } - - // ── Num images mapping ────────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - assert_plan_for_batch(None, |plan| { - assert!(matches!(plan.num_images, FalFlux1DevNumImages::One)); - }); - } - - #[test] - fn batch_of_one_yields_one() { - assert_plan_for_batch(Some(1), |plan| { - assert!(matches!(plan.num_images, FalFlux1DevNumImages::One)); - }); - } - - #[test] - fn batch_of_two_yields_two() { - assert_plan_for_batch(Some(2), |plan| { - assert!(matches!(plan.num_images, FalFlux1DevNumImages::Two)); - }); - } - - #[test] - fn batch_of_three_yields_three() { - assert_plan_for_batch(Some(3), |plan| { - assert!(matches!(plan.num_images, FalFlux1DevNumImages::Three)); - }); - } - - #[test] - fn batch_of_four_yields_four() { - assert_plan_for_batch(Some(4), |plan| { - assert!(matches!(plan.num_images, FalFlux1DevNumImages::Four)); - }); - } - - #[test] - fn batch_above_four_clamps_to_four() { - // PayMoreUpgrade strategy clamps over-large batches down to the max. - assert_plan_for_batch(Some(9), |plan| { - assert!(matches!(plan.num_images, FalFlux1DevNumImages::Four)); - }); - } - - // ── Prompt passthrough ────────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_fal_plan(&make_request(Some("a corgi in a hat"), None, Some(1)), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_fal_plan(&make_request(None, None, Some(1)), |plan| { - assert_eq!(plan.prompt, None); - }); - } - - // ── Image refs (Flux 1 Dev never has them) ────────────────────────────── - - #[test] - fn image_url_is_always_none() { - // Flux 1 Dev is text-to-image only — `maybe_image_url` is always None - // because we never pass image_media_tokens. - with_fal_plan(&make_request(Some("p"), None, Some(1)), |plan| { - assert!(plan.maybe_image_url.is_none()); - }); - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_schnell.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_schnell.rs deleted file mode 100644 index 80fb1b96cb..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_schnell.rs +++ /dev/null @@ -1,529 +0,0 @@ -//! Tests for omni-gen against the Flux 1 Schnell model. -//! -//! These exercise [`distill_image_request`] end-to-end without firing any -//! provider calls — both the cost estimate (Artcraft provider, what we'd -//! bill on) and the execution plan (Fal provider, what we'd send to fal.ai) -//! are computed synchronously from the request. -//! -//! Flux 1 Schnell is text-to-image only, so no image refs are exercised. -//! It's also currently free / unlimited — we don't charge anything for it. - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - /// Build a minimal Flux 1 Schnell omni request with only the fields we care - /// about in tests filled in. - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - image_batch_count: Option, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::Flux1Schnell), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens: None, - resolution: None, - aspect_ratio, - quality: None, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - /// Distill a request and panic with a descriptive message on failure. - /// Flux 1 Schnell never uses image refs, so we always pass an empty hydration map. - fn distill(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for Flux 1 Schnell") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST - // ──────────────────────────────────────────────────────────────────────────── - // - // Pricing comes from `estimate_image_cost_artcraft_flux_1_schnell`: - // Flux 1 Schnell is currently FREE / UNLIMITED — we don't charge for it. - // - cost_in_credits: None - // - cost_in_usd_cents: None - // - is_free: true - // - is_unlimited: true - // - // The artcraft cost path is what omni would bill on. - mod cost { - use super::*; - - fn distilled_for(image_batch_count: Option) -> DistilledImageRequest { - distill(&make_request(Some("a happy puppy"), None, image_batch_count)) - } - - #[test] - fn default_batch_is_free() { - let distilled = distilled_for(None); - assert_eq!(distilled.cost.cost_in_credits, None); - assert_eq!(distilled.cost.cost_in_usd_cents, None); - assert!(distilled.cost.is_free); - assert!(distilled.cost.is_unlimited); - } - - #[test] - fn batch_of_one_is_free() { - let distilled = distilled_for(Some(1)); - assert_eq!(distilled.cost.cost_in_credits, None); - assert_eq!(distilled.cost.cost_in_usd_cents, None); - assert!(distilled.cost.is_free); - assert!(distilled.cost.is_unlimited); - } - - #[test] - fn batch_of_two_is_free() { - let distilled = distilled_for(Some(2)); - assert_eq!(distilled.cost.cost_in_credits, None); - assert_eq!(distilled.cost.cost_in_usd_cents, None); - assert!(distilled.cost.is_free); - assert!(distilled.cost.is_unlimited); - } - - #[test] - fn batch_of_three_is_free() { - let distilled = distilled_for(Some(3)); - assert_eq!(distilled.cost.cost_in_credits, None); - assert_eq!(distilled.cost.cost_in_usd_cents, None); - assert!(distilled.cost.is_free); - assert!(distilled.cost.is_unlimited); - } - - #[test] - fn batch_of_four_is_free() { - let distilled = distilled_for(Some(4)); - assert_eq!(distilled.cost.cost_in_credits, None); - assert_eq!(distilled.cost.cost_in_usd_cents, None); - assert!(distilled.cost.is_free); - assert!(distilled.cost.is_unlimited); - } - - #[test] - fn batch_above_max_is_still_free() { - // PayMoreUpgrade strategy clamps batches >4 down to the max of 4 — and - // 4-of-free is still free. - let distilled = distilled_for(Some(7)); - assert_eq!(distilled.cost.cost_in_credits, None); - assert_eq!(distilled.cost.cost_in_usd_cents, None); - assert!(distilled.cost.is_free); - assert!(distilled.cost.is_unlimited); - } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let aspect_ratios = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::SquareHd), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::TallNineBySixteen), - Some(CommonAspectRatio::Auto), - ]; - for ar in aspect_ratios { - let request = make_request(Some("p"), ar, Some(2)); - let distilled = distill(&request); - assert_eq!( - distilled.cost.cost_in_usd_cents, None, - "expected free regardless of aspect ratio (got {:?})", ar, - ); - assert!(distilled.cost.is_free); - assert!(distilled.cost.is_unlimited); - } - } - - #[test] - fn cost_is_independent_of_prompt() { - let with_prompt = distill(&make_request(Some("a cat"), None, Some(3))); - let without_prompt = distill(&make_request(None, None, Some(3))); - - assert_eq!(with_prompt.cost.cost_in_usd_cents, None); - assert_eq!(without_prompt.cost.cost_in_usd_cents, None); - assert!(with_prompt.cost.is_free); - assert!(without_prompt.cost.is_free); - } - - #[test] - fn cost_metadata_flags() { - let request = make_request(Some("p"), None, Some(1)); - let distilled = distill(&request); - assert!(distilled.cost.is_free); - assert!(distilled.cost.is_unlimited); - // Free / unlimited but not rate-limited and no watermark. - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - // - // The distilled plan is the Fal-side `PlanFalFlux1Schnell`: - // - prompt: passthrough - // - maybe_image_url: always None (Flux 1 Schnell never gets image refs in tests) - // - t2i_aspect_ratio: text-to-image aspect ratio (always set; defaults to Square) - // - edit_image_size: image-edit size (only populated when caller passes an - // aspect ratio; None when no aspect ratio was supplied) - // - num_images: mapped from image_batch_count (default 1, clamped to 1..=4) - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_flux_1_schnell::{ - FalFlux1SchnellNumImages, PlanFalFlux1Schnell, - }; - use fal_client::requests::webhook::image::edit::enqueue_flux_1_schnell_edit_image_webhook::Flux1SchnellEditImageSize; - use fal_client::requests::webhook::image::text::enqueue_flux_1_schnell_text_to_image_webhook::Flux1SchnellAspectRatio; - - /// Distill and run an assertion against the inner FalFlux1Schnell plan. - /// - /// We pass the plan by reference (rather than moving it out of `distilled`) - /// because the plan's borrows are tied to data owned by `distilled` — - /// moving it out would dangle those borrows. - fn with_fal_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill(request); - match distilled.plan() { - ImageGenerationPlan::FalFlux1Schnell(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalFlux1Schnell, got {:?}", other), - } - } - - fn assert_plan_for_aspect_ratio( - ar: Option, - assertion: F, - ) { - with_fal_plan(&make_request(Some("p"), ar, Some(1)), assertion); - } - - fn assert_plan_for_batch( - batch: Option, - assertion: F, - ) { - with_fal_plan(&make_request(Some("p"), None, batch), assertion); - } - - // ── t2i_aspect_ratio direct mapping ───────────────────────────────────── - - #[test] - fn default_t2i_aspect_ratio_is_square() { - assert_plan_for_aspect_ratio(None, |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::Square)); - }); - } - - #[test] - fn square_input_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::Square)); - }); - } - - #[test] - fn square_hd_input_yields_square_hd() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::SquareHd)); - }); - } - - #[test] - fn wide_4x3_yields_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn wide_16x9_yields_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn wide_alias_yields_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn tall_3x4_yields_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::PortraitThreeByFour)); - }); - } - - #[test] - fn tall_9x16_yields_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::PortraitNineBySixteen)); - }); - } - - #[test] - fn tall_alias_yields_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::PortraitNineBySixteen)); - }); - } - - // ── t2i_aspect_ratio nearest-neighbour fallbacks ──────────────────────── - // - // Aspect ratios that Flux 1 Schnell doesn't natively support get mapped to - // the closest supported ratio under the PayMoreUpgrade mismatch strategy - // that omni-gen uses. - - #[test] - fn wide_5x4_falls_back_to_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn wide_3x2_falls_back_to_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideThreeByTwo), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn wide_21x9_falls_back_to_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideTwentyOneByNine), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn tall_4x5_falls_back_to_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallFourByFive), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::PortraitThreeByFour)); - }); - } - - #[test] - fn tall_2x3_falls_back_to_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::PortraitThreeByFour)); - }); - } - - #[test] - fn tall_9x21_falls_back_to_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallNineByTwentyOne), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::PortraitNineBySixteen)); - }); - } - - // ── Auto* aspect ratios collapse to Square (text-to-image) ────────────── - - #[test] - fn auto_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::Square)); - }); - } - - #[test] - fn auto_2k_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto2k), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::Square)); - }); - } - - #[test] - fn auto_4k_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto4k), |plan| { - assert!(matches!(plan.t2i_aspect_ratio, Flux1SchnellAspectRatio::Square)); - }); - } - - // ── edit_image_size mapping ───────────────────────────────────────────── - // - // The plan also computes an `edit_image_size` from the same aspect ratio - // input. It's None when no aspect ratio was supplied, and otherwise mirrors - // the t2i mapping (with a Square fallback for unsupported variants). - - #[test] - fn no_aspect_ratio_yields_no_edit_image_size() { - assert_plan_for_aspect_ratio(None, |plan| { - assert!(plan.edit_image_size.is_none()); - }); - } - - #[test] - fn square_yields_edit_image_size_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::Square))); - }); - } - - #[test] - fn square_hd_yields_edit_image_size_square_hd() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::SquareHd))); - }); - } - - #[test] - fn wide_4x3_yields_edit_image_size_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::LandscapeFourByThree))); - }); - } - - #[test] - fn wide_16x9_yields_edit_image_size_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::LandscapeSixteenByNine))); - }); - } - - #[test] - fn wide_alias_yields_edit_image_size_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::LandscapeSixteenByNine))); - }); - } - - #[test] - fn tall_3x4_yields_edit_image_size_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::PortraitThreeByFour))); - }); - } - - #[test] - fn tall_9x16_yields_edit_image_size_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::PortraitNineBySixteen))); - }); - } - - #[test] - fn tall_alias_yields_edit_image_size_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::PortraitNineBySixteen))); - }); - } - - /// Aspect ratios that don't have a direct edit-image-size mapping (Auto, - /// the Wide5x4 / Wide3x2 / Wide21x9 family, the Tall4x5 / Tall2x3 / - /// Tall9x21 family) all collapse to `Some(Square)` per the plan's catch-all. - #[test] - fn auto_yields_edit_image_size_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::Square))); - }); - } - - #[test] - fn wide_5x4_yields_edit_image_size_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::Square))); - }); - } - - #[test] - fn tall_2x3_yields_edit_image_size_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!(plan.edit_image_size, Some(Flux1SchnellEditImageSize::Square))); - }); - } - - // ── Num images mapping ────────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - assert_plan_for_batch(None, |plan| { - assert!(matches!(plan.num_images, FalFlux1SchnellNumImages::One)); - }); - } - - #[test] - fn batch_of_one_yields_one() { - assert_plan_for_batch(Some(1), |plan| { - assert!(matches!(plan.num_images, FalFlux1SchnellNumImages::One)); - }); - } - - #[test] - fn batch_of_two_yields_two() { - assert_plan_for_batch(Some(2), |plan| { - assert!(matches!(plan.num_images, FalFlux1SchnellNumImages::Two)); - }); - } - - #[test] - fn batch_of_three_yields_three() { - assert_plan_for_batch(Some(3), |plan| { - assert!(matches!(plan.num_images, FalFlux1SchnellNumImages::Three)); - }); - } - - #[test] - fn batch_of_four_yields_four() { - assert_plan_for_batch(Some(4), |plan| { - assert!(matches!(plan.num_images, FalFlux1SchnellNumImages::Four)); - }); - } - - #[test] - fn batch_above_four_clamps_to_four() { - // PayMoreUpgrade strategy clamps over-large batches down to the max. - assert_plan_for_batch(Some(9), |plan| { - assert!(matches!(plan.num_images, FalFlux1SchnellNumImages::Four)); - }); - } - - // ── Prompt passthrough ────────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_fal_plan(&make_request(Some("a corgi in a hat"), None, Some(1)), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_fal_plan(&make_request(None, None, Some(1)), |plan| { - assert_eq!(plan.prompt, None); - }); - } - - // ── Image refs (Flux 1 Schnell never has them) ────────────────────────── - - #[test] - fn image_url_is_always_none() { - // Flux 1 Schnell is text-to-image only — `maybe_image_url` is always - // None because we never pass image_media_tokens. - with_fal_plan(&make_request(Some("p"), None, Some(1)), |plan| { - assert!(plan.maybe_image_url.is_none()); - }); - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1.rs deleted file mode 100644 index be69f5e709..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1.rs +++ /dev/null @@ -1,412 +0,0 @@ -//! Tests for omni-gen against the Flux Pro 1.1 model. -//! -//! These exercise [`distill_image_request`] end-to-end without firing any -//! provider calls — both the cost estimate (Artcraft provider, what we bill -//! on) and the execution plan (Fal provider, what we'd send to fal.ai) are -//! computed synchronously from the request. -//! -//! Flux Pro 1.1 is text-to-image only, so no image refs are exercised. - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - /// Build a minimal Flux Pro 1.1 omni request with only the fields we care - /// about in tests filled in. - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - image_batch_count: Option, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::FluxPro11), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens: None, - resolution: None, - aspect_ratio, - quality: None, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - /// Distill a request and panic with a descriptive message on failure. - /// Flux Pro 1.1 never uses image refs, so we always pass an empty hydration map. - fn distill(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for Flux Pro 1.1") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST - // ──────────────────────────────────────────────────────────────────────────── - // - // Pricing comes from `estimate_image_cost_artcraft_flux_pro_1p1`: - // $0.04 / image (4 credits per image; 1 credit == 1 USD cent) - // - // This matches what the legacy `generate_flux_pro_11_text_to_image_handler` - // charges via `FluxPro11Args::calculate_cost_in_cents` (which is also $0.04 - // per image — Fal docs: "$0.04 per megapixel, billed by rounding up to the - // nearest megapixel" at the default ~1MP image size). - // - // The artcraft cost path is what omni bills on. - mod cost { - use super::*; - - fn cost_for_batch(image_batch_count: Option) -> (Option, Option) { - let request = make_request(Some("a happy puppy"), None, image_batch_count); - let distilled = distill(&request); - ( - distilled.cost.cost_in_credits, - distilled.cost.cost_in_usd_cents, - ) - } - - #[test] - fn default_batch_costs_4_cents() { - // No batch specified → defaults to 1 image → 4¢ - let (credits, cents) = cost_for_batch(None); - assert_eq!(credits, Some(4)); - assert_eq!(cents, Some(4)); - } - - #[test] - fn batch_of_one_costs_4_cents() { - let (credits, cents) = cost_for_batch(Some(1)); - assert_eq!(credits, Some(4)); - assert_eq!(cents, Some(4)); - } - - #[test] - fn batch_of_two_costs_8_cents() { - let (credits, cents) = cost_for_batch(Some(2)); - assert_eq!(credits, Some(8)); - assert_eq!(cents, Some(8)); - } - - #[test] - fn batch_of_three_costs_12_cents() { - let (credits, cents) = cost_for_batch(Some(3)); - assert_eq!(credits, Some(12)); - assert_eq!(cents, Some(12)); - } - - #[test] - fn batch_of_four_costs_16_cents() { - let (credits, cents) = cost_for_batch(Some(4)); - assert_eq!(credits, Some(16)); - assert_eq!(cents, Some(16)); - } - - #[test] - fn batch_above_max_clamps_and_costs_16_cents() { - // PayMoreUpgrade strategy clamps batches >4 down to the max of 4. - let (credits, cents) = cost_for_batch(Some(7)); - assert_eq!(credits, Some(16)); - assert_eq!(cents, Some(16)); - } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let aspect_ratios = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::SquareHd), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::TallNineBySixteen), - Some(CommonAspectRatio::Auto), - ]; - for ar in aspect_ratios { - let request = make_request(Some("p"), ar, Some(2)); - let distilled = distill(&request); - assert_eq!( - distilled.cost.cost_in_usd_cents, - Some(8), - "expected 8¢ regardless of aspect ratio (got {:?})", - ar, - ); - } - } - - #[test] - fn cost_is_independent_of_prompt() { - let with_prompt = distill(&make_request(Some("a cat"), None, Some(3))); - let without_prompt = distill(&make_request(None, None, Some(3))); - assert_eq!(with_prompt.cost.cost_in_usd_cents, Some(12)); - assert_eq!(without_prompt.cost.cost_in_usd_cents, Some(12)); - } - - #[test] - fn cost_metadata_flags_are_default() { - let request = make_request(Some("p"), None, Some(1)); - let distilled = distill(&request); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - // - // The distilled plan is the Fal-side `PlanFalFluxPro11`: - // - prompt: passthrough - // - aspect_ratio: mapped from CommonAspectRatio with sensible fallbacks - // - num_images: mapped from image_batch_count (default 1, clamped to 1..=4) - // - // (Flux Pro 1.1 has no image_url field — it's text-to-image only.) - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_flux_pro_1p1::{ - FalFluxPro11NumImages, PlanFalFluxPro11, - }; - use fal_client::requests::webhook::image::text::enqueue_flux_pro_11_text_to_image_webhook::FluxPro11AspectRatio; - - /// Distill and run an assertion against the inner FalFluxPro11 plan. - /// - /// We pass the plan by reference (rather than moving it out of `distilled`) - /// because the plan's borrows are tied to data owned by `distilled` — - /// moving it out would dangle those borrows. - fn with_fal_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill(request); - match distilled.plan() { - ImageGenerationPlan::FalFluxPro11(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalFluxPro11, got {:?}", other), - } - } - - fn assert_plan_for_aspect_ratio( - ar: Option, - assertion: F, - ) { - with_fal_plan(&make_request(Some("p"), ar, Some(1)), assertion); - } - - fn assert_plan_for_batch( - batch: Option, - assertion: F, - ) { - with_fal_plan(&make_request(Some("p"), None, batch), assertion); - } - - // ── Aspect ratio direct mappings ──────────────────────────────────────── - - #[test] - fn default_aspect_ratio_is_square() { - assert_plan_for_aspect_ratio(None, |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::Square)); - }); - } - - #[test] - fn square_input_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::Square)); - }); - } - - #[test] - fn square_hd_input_yields_square_hd() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::SquareHd)); - }); - } - - #[test] - fn wide_4x3_yields_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn wide_16x9_yields_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn wide_alias_yields_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn tall_3x4_yields_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::PortraitThreeByFour)); - }); - } - - #[test] - fn tall_9x16_yields_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::PortraitNineBySixteen)); - }); - } - - #[test] - fn tall_alias_yields_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::PortraitNineBySixteen)); - }); - } - - // ── Aspect ratio nearest-neighbour fallbacks ──────────────────────────── - // - // Aspect ratios that Flux Pro 1.1 doesn't natively support get mapped to - // the closest supported ratio under the PayMoreUpgrade mismatch strategy - // that omni-gen uses. - - #[test] - fn wide_5x4_falls_back_to_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn wide_3x2_falls_back_to_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideThreeByTwo), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn wide_21x9_falls_back_to_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideTwentyOneByNine), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn tall_4x5_falls_back_to_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallFourByFive), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::PortraitThreeByFour)); - }); - } - - #[test] - fn tall_2x3_falls_back_to_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::PortraitThreeByFour)); - }); - } - - #[test] - fn tall_9x21_falls_back_to_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallNineByTwentyOne), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::PortraitNineBySixteen)); - }); - } - - // ── Auto* aspect ratios ───────────────────────────────────────────────── - - #[test] - fn auto_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::Square)); - }); - } - - #[test] - fn auto_2k_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto2k), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::Square)); - }); - } - - #[test] - fn auto_4k_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto4k), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11AspectRatio::Square)); - }); - } - - // ── Num images mapping ────────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - assert_plan_for_batch(None, |plan| { - assert!(matches!(plan.num_images, FalFluxPro11NumImages::One)); - }); - } - - #[test] - fn batch_of_one_yields_one() { - assert_plan_for_batch(Some(1), |plan| { - assert!(matches!(plan.num_images, FalFluxPro11NumImages::One)); - }); - } - - #[test] - fn batch_of_two_yields_two() { - assert_plan_for_batch(Some(2), |plan| { - assert!(matches!(plan.num_images, FalFluxPro11NumImages::Two)); - }); - } - - #[test] - fn batch_of_three_yields_three() { - assert_plan_for_batch(Some(3), |plan| { - assert!(matches!(plan.num_images, FalFluxPro11NumImages::Three)); - }); - } - - #[test] - fn batch_of_four_yields_four() { - assert_plan_for_batch(Some(4), |plan| { - assert!(matches!(plan.num_images, FalFluxPro11NumImages::Four)); - }); - } - - #[test] - fn batch_above_four_clamps_to_four() { - // PayMoreUpgrade strategy clamps over-large batches down to the max. - assert_plan_for_batch(Some(9), |plan| { - assert!(matches!(plan.num_images, FalFluxPro11NumImages::Four)); - }); - } - - // ── Prompt passthrough ────────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_fal_plan(&make_request(Some("a corgi in a hat"), None, Some(1)), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_fal_plan(&make_request(None, None, Some(1)), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1_ultra.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1_ultra.rs deleted file mode 100644 index 666f7974e9..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1_ultra.rs +++ /dev/null @@ -1,417 +0,0 @@ -//! Tests for omni-gen against the Flux Pro 1.1 Ultra model. -//! -//! These exercise [`distill_image_request`] end-to-end without firing any -//! provider calls — both the cost estimate (Artcraft provider, what we bill -//! on) and the execution plan (Fal provider, what we'd send to fal.ai) are -//! computed synchronously from the request. -//! -//! Flux Pro 1.1 Ultra is text-to-image only, so no image refs are exercised. - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - /// Build a minimal Flux Pro 1.1 Ultra omni request with only the fields we - /// care about in tests filled in. - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - image_batch_count: Option, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::FluxPro11Ultra), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens: None, - resolution: None, - aspect_ratio, - quality: None, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - /// Distill a request and panic with a descriptive message on failure. - /// Flux Pro 1.1 Ultra never uses image refs, so we always pass an empty - /// hydration map. - fn distill(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for Flux Pro 1.1 Ultra") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST - // ──────────────────────────────────────────────────────────────────────────── - // - // Pricing comes from `estimate_image_cost_artcraft_flux_pro_1p1_ultra`: - // $0.06 / image (6 credits per image; 1 credit == 1 USD cent) - // - // This matches what the legacy - // `generate_flux_pro_11_ultra_text_to_image_handler` charges via - // `FluxPro11UltraArgs::calculate_cost_in_cents` (Fal docs: - // "$0.06 per image", flat). - mod cost { - use super::*; - - fn cost_for_batch(image_batch_count: Option) -> (Option, Option) { - let request = make_request(Some("a happy puppy"), None, image_batch_count); - let distilled = distill(&request); - ( - distilled.cost.cost_in_credits, - distilled.cost.cost_in_usd_cents, - ) - } - - #[test] - fn default_batch_costs_6_cents() { - // No batch specified → defaults to 1 image → 6¢ - let (credits, cents) = cost_for_batch(None); - assert_eq!(credits, Some(6)); - assert_eq!(cents, Some(6)); - } - - #[test] - fn batch_of_one_costs_6_cents() { - let (credits, cents) = cost_for_batch(Some(1)); - assert_eq!(credits, Some(6)); - assert_eq!(cents, Some(6)); - } - - #[test] - fn batch_of_two_costs_12_cents() { - let (credits, cents) = cost_for_batch(Some(2)); - assert_eq!(credits, Some(12)); - assert_eq!(cents, Some(12)); - } - - #[test] - fn batch_of_three_costs_18_cents() { - let (credits, cents) = cost_for_batch(Some(3)); - assert_eq!(credits, Some(18)); - assert_eq!(cents, Some(18)); - } - - #[test] - fn batch_of_four_costs_24_cents() { - let (credits, cents) = cost_for_batch(Some(4)); - assert_eq!(credits, Some(24)); - assert_eq!(cents, Some(24)); - } - - #[test] - fn batch_above_max_clamps_and_costs_24_cents() { - // PayMoreUpgrade strategy clamps batches >4 down to the max of 4. - let (credits, cents) = cost_for_batch(Some(7)); - assert_eq!(credits, Some(24)); - assert_eq!(cents, Some(24)); - } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let aspect_ratios = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::SquareHd), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::WideTwentyOneByNine), - Some(CommonAspectRatio::TallNineBySixteen), - Some(CommonAspectRatio::TallNineByTwentyOne), - Some(CommonAspectRatio::Auto), - ]; - for ar in aspect_ratios { - let request = make_request(Some("p"), ar, Some(2)); - let distilled = distill(&request); - assert_eq!( - distilled.cost.cost_in_usd_cents, - Some(12), - "expected 12¢ regardless of aspect ratio (got {:?})", - ar, - ); - } - } - - #[test] - fn cost_is_independent_of_prompt() { - let with_prompt = distill(&make_request(Some("a cat"), None, Some(3))); - let without_prompt = distill(&make_request(None, None, Some(3))); - assert_eq!(with_prompt.cost.cost_in_usd_cents, Some(18)); - assert_eq!(without_prompt.cost.cost_in_usd_cents, Some(18)); - } - - #[test] - fn cost_metadata_flags_are_default() { - let request = make_request(Some("p"), None, Some(1)); - let distilled = distill(&request); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - // - // The distilled plan is the Fal-side `PlanFalFluxPro11Ultra`: - // - prompt: passthrough - // - aspect_ratio: mapped from CommonAspectRatio (Ultra has more native - // ratios than 1.1, including 21:9, 3:2, 9:21, 2:3) - // - num_images: mapped from image_batch_count (default 1, clamped to 1..=4) - // - // (Flux Pro 1.1 Ultra has no image_url field — text-to-image only.) - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_flux_pro_1p1_ultra::{ - FalFluxPro11UltraNumImages, PlanFalFluxPro11Ultra, - }; - use fal_client::requests::webhook::image::text::enqueue_flux_pro_11_ultra_text_to_image_webhook::FluxPro11UltraAspectRatio; - - /// Distill and run an assertion against the inner FalFluxPro11Ultra plan. - /// - /// We pass the plan by reference (rather than moving it out of `distilled`) - /// because the plan's borrows are tied to data owned by `distilled` — - /// moving it out would dangle those borrows. - fn with_fal_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill(request); - match distilled.plan() { - ImageGenerationPlan::FalFluxPro11Ultra(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalFluxPro11Ultra, got {:?}", other), - } - } - - fn assert_plan_for_aspect_ratio( - ar: Option, - assertion: F, - ) { - with_fal_plan(&make_request(Some("p"), ar, Some(1)), assertion); - } - - fn assert_plan_for_batch( - batch: Option, - assertion: F, - ) { - with_fal_plan(&make_request(Some("p"), None, batch), assertion); - } - - // ── Aspect ratio direct mappings ──────────────────────────────────────── - - #[test] - fn default_aspect_ratio_is_square() { - assert_plan_for_aspect_ratio(None, |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::Square)); - }); - } - - #[test] - fn square_input_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::Square)); - }); - } - - #[test] - fn square_hd_yields_square() { - // Ultra collapses SquareHd back to plain Square. - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::Square)); - }); - } - - #[test] - fn wide_4x3_yields_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn wide_3x2_yields_landscape_3_2() { - // Ultra-only direct mapping (1.1 doesn't have this). - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideThreeByTwo), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::LandscapeThreeByTwo)); - }); - } - - #[test] - fn wide_16x9_yields_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn wide_alias_yields_landscape_16_9() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::LandscapeSixteenByNine)); - }); - } - - #[test] - fn wide_21x9_yields_landscape_21_9() { - // Ultra-only direct mapping (1.1 falls back to 16:9). - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideTwentyOneByNine), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::LandscapeTwentyOneByNine)); - }); - } - - #[test] - fn tall_3x4_yields_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::PortraitThreeByFour)); - }); - } - - #[test] - fn tall_2x3_yields_portrait_2_3() { - // Ultra-only direct mapping (1.1 falls back to 3:4). - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::PortraitTwoByThree)); - }); - } - - #[test] - fn tall_9x16_yields_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::PortraitNineBySixteen)); - }); - } - - #[test] - fn tall_alias_yields_portrait_9_16() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::PortraitNineBySixteen)); - }); - } - - #[test] - fn tall_9x21_yields_portrait_9_21() { - // Ultra-only direct mapping (1.1 falls back to 9:16). - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallNineByTwentyOne), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::PortraitNineByTwentyOne)); - }); - } - - // ── Aspect ratio nearest-neighbour fallbacks ──────────────────────────── - // - // Ultra natively supports almost everything; only 5:4 and 4:5 fall back. - - #[test] - fn wide_5x4_falls_back_to_landscape_4_3() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::LandscapeFourByThree)); - }); - } - - #[test] - fn tall_4x5_falls_back_to_portrait_3_4() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::TallFourByFive), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::PortraitThreeByFour)); - }); - } - - // ── Auto* aspect ratios ───────────────────────────────────────────────── - - #[test] - fn auto_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::Square)); - }); - } - - #[test] - fn auto_2k_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto2k), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::Square)); - }); - } - - #[test] - fn auto_4k_yields_square() { - assert_plan_for_aspect_ratio(Some(CommonAspectRatio::Auto4k), |plan| { - assert!(matches!(plan.aspect_ratio, FluxPro11UltraAspectRatio::Square)); - }); - } - - // ── Num images mapping ────────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - assert_plan_for_batch(None, |plan| { - assert!(matches!(plan.num_images, FalFluxPro11UltraNumImages::One)); - }); - } - - #[test] - fn batch_of_one_yields_one() { - assert_plan_for_batch(Some(1), |plan| { - assert!(matches!(plan.num_images, FalFluxPro11UltraNumImages::One)); - }); - } - - #[test] - fn batch_of_two_yields_two() { - assert_plan_for_batch(Some(2), |plan| { - assert!(matches!(plan.num_images, FalFluxPro11UltraNumImages::Two)); - }); - } - - #[test] - fn batch_of_three_yields_three() { - assert_plan_for_batch(Some(3), |plan| { - assert!(matches!(plan.num_images, FalFluxPro11UltraNumImages::Three)); - }); - } - - #[test] - fn batch_of_four_yields_four() { - assert_plan_for_batch(Some(4), |plan| { - assert!(matches!(plan.num_images, FalFluxPro11UltraNumImages::Four)); - }); - } - - #[test] - fn batch_above_four_clamps_to_four() { - // PayMoreUpgrade strategy clamps over-large batches down to the max. - assert_plan_for_batch(Some(9), |plan| { - assert!(matches!(plan.num_images, FalFluxPro11UltraNumImages::Four)); - }); - } - - // ── Prompt passthrough ────────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_fal_plan(&make_request(Some("a corgi in a hat"), None, Some(1)), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_fal_plan(&make_request(None, None, Some(1)), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1.rs deleted file mode 100644 index bc9875d44e..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1.rs +++ /dev/null @@ -1,823 +0,0 @@ -//! Tests for omni-gen against the GPT Image 1 model. -//! -//! GPT Image 1 is multi-function: text-to-image and image-to-image (edit). -//! The omni-gen distillation routes through Fal for the execution plan and -//! Artcraft for billing. -//! -//! Artcraft-tier pricing (the rate the user is billed): -//! -//! Output image cost (per output image): -//! Low – 2¢ (square/unset) / 2¢ (wide or tall) -//! Medium – 5¢ (square/unset) / 7¢ (wide or tall) -//! High – 17¢ (square/unset) / 25¢ (wide or tall) -//! None → defaults to High -//! -//! Input image cost (edit mode only, per input image): -//! 2¢ per input image (high-fidelity token estimate) -//! -//! The Fal execution plan maps quality from the omni request (defaulting to -//! Medium when unspecified) and maps aspect ratios to three image sizes: -//! Square (1024×1024), Horizontal (1536×1024), Vertical (1024×1536). -//! Auto/unset maps to None (API default). - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use enums::common::generation::common_quality::CommonQuality; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - quality: Option, - image_batch_count: Option, - image_media_tokens: Option>, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::GptImage1), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens, - resolution: None, - aspect_ratio, - quality, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - fn fake_image_refs(n: usize) -> (Vec, HashMap) { - let mut tokens = Vec::with_capacity(n); - let mut map = HashMap::with_capacity(n); - for i in 0..n { - let token = MediaFileToken::new_from_str(&format!("mf_test{:028}", i)); - let url = Url::parse(&format!("https://fake.example.com/img{}.png", i)).unwrap(); - map.insert(token.clone(), url); - tokens.push(token); - } - (tokens, map) - } - - fn distill_text(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for GPT Image 1 (text)") - } - - fn distill_edit( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - ) -> DistilledImageRequest { - distill_image_request(request, Some(hydration)) - .expect("distill_image_request should succeed for GPT Image 1 (edit)") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST - // ──────────────────────────────────────────────────────────────────────────── - mod cost { - use super::*; - - mod text { - use super::*; - - fn cost( - quality: Option, - aspect_ratio: Option, - batch: Option, - ) -> u64 { - let request = make_request(Some("p"), aspect_ratio, quality, batch, None); - distill_text(&request).cost.cost_in_usd_cents.unwrap() - } - - // ── Default quality (None → High) ──────────────────────────────────── - - #[test] - fn default_quality_square_costs_17() { - assert_eq!(cost(None, Some(CommonAspectRatio::Square), Some(1)), 17); - } - - #[test] - fn default_quality_wide_costs_25() { - assert_eq!(cost(None, Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 25); - } - - #[test] - fn default_quality_tall_costs_25() { - assert_eq!(cost(None, Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 25); - } - - #[test] - fn default_quality_unset_size_costs_17() { - // Unset aspect_ratio → None image_size → treated as square = 17¢ - assert_eq!(cost(None, None, Some(1)), 17); - } - - #[test] - fn default_quality_auto_costs_17() { - assert_eq!(cost(None, Some(CommonAspectRatio::Auto), Some(1)), 17); - } - - // ── Low quality (2¢/image regardless of size) ──────────────────────── - - #[test] - fn low_square_costs_2() { - assert_eq!(cost(Some(CommonQuality::Low), Some(CommonAspectRatio::Square), Some(1)), 2); - } - - #[test] - fn low_wide_costs_2() { - assert_eq!(cost(Some(CommonQuality::Low), Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 2); - } - - #[test] - fn low_tall_costs_2() { - assert_eq!(cost(Some(CommonQuality::Low), Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 2); - } - - #[test] - fn low_unset_costs_2() { - assert_eq!(cost(Some(CommonQuality::Low), None, Some(1)), 2); - } - - #[test] - fn low_four_images_costs_8() { - assert_eq!(cost(Some(CommonQuality::Low), None, Some(4)), 8); - } - - #[test] - fn low_batch_above_max_clamps_to_8() { - assert_eq!(cost(Some(CommonQuality::Low), None, Some(7)), 8); - } - - // ── Medium quality (5¢ square, 7¢ wide/tall) ───────────────────────── - - #[test] - fn medium_square_costs_5() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::Square), Some(1)), 5); - } - - #[test] - fn medium_wide_costs_7() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 7); - } - - #[test] - fn medium_tall_costs_7() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 7); - } - - #[test] - fn medium_unset_costs_5() { - assert_eq!(cost(Some(CommonQuality::Medium), None, Some(1)), 5); - } - - #[test] - fn medium_square_four_images_costs_20() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::Square), Some(4)), 20); - } - - #[test] - fn medium_wide_four_images_costs_28() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::WideSixteenByNine), Some(4)), 28); - } - - // ── High quality (17¢ square, 25¢ wide/tall) ───────────────────────── - - #[test] - fn high_square_costs_17() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1)), 17); - } - - #[test] - fn high_wide_costs_25() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 25); - } - - #[test] - fn high_tall_costs_25() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 25); - } - - #[test] - fn high_unset_costs_17() { - assert_eq!(cost(Some(CommonQuality::High), None, Some(1)), 17); - } - - #[test] - fn high_square_four_images_costs_68() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(4)), 68); - } - - #[test] - fn high_wide_four_images_costs_100() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::WideSixteenByNine), Some(4)), 100); - } - - #[test] - fn high_batch_above_max_clamps_to_four() { - assert_eq!(cost(Some(CommonQuality::High), None, Some(9)), 68); - } - - // ── Batch scaling ──────────────────────────────────────────────────── - - #[test] - fn high_square_batch_scaling() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1)), 17); - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(2)), 34); - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(3)), 51); - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(4)), 68); - } - - // ── Metadata flags ─────────────────────────────────────────────────── - - #[test] - fn cost_metadata_flags_are_default() { - let request = make_request(Some("p"), None, None, Some(1), None); - let distilled = distill_text(&request); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - - // ── Cost is independent of prompt ───────────────────────────────────── - - #[test] - fn cost_is_independent_of_prompt() { - let with_prompt = distill_text(&make_request(Some("a cat"), None, None, Some(1), None)); - let without_prompt = distill_text(&make_request(None, None, None, Some(1), None)); - assert_eq!(with_prompt.cost.cost_in_usd_cents, without_prompt.cost.cost_in_usd_cents); - } - } - - mod edit { - use super::*; - - fn cost_edit( - quality: Option, - aspect_ratio: Option, - batch: Option, - num_input_images: usize, - ) -> u64 { - let (tokens, hydration) = fake_image_refs(num_input_images); - let request = make_request(Some("p"), aspect_ratio, quality, batch, Some(tokens)); - distill_edit(&request, &hydration).cost.cost_in_usd_cents.unwrap() - } - - // ── Input image cost adds 2¢ per input ────────────────────────────── - - #[test] - fn edit_high_square_one_output_one_input() { - // 17¢ output + 1×2¢ input = 19¢ - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1), 1), 19); - } - - #[test] - fn edit_high_square_one_output_three_inputs() { - // 17¢ output + 3×2¢ input = 23¢ - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1), 3), 23); - } - - #[test] - fn edit_high_square_one_output_five_inputs() { - // 17¢ output + 5×2¢ input = 27¢ - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1), 5), 27); - } - - #[test] - fn edit_high_wide_two_outputs_one_input() { - // 2×25¢ output + 1×2¢ input = 52¢ - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::WideSixteenByNine), Some(2), 1), 52); - } - - #[test] - fn edit_medium_wide_two_outputs_five_inputs() { - // 2×7¢ output + 5×2¢ input = 24¢ - assert_eq!(cost_edit(Some(CommonQuality::Medium), Some(CommonAspectRatio::WideSixteenByNine), Some(2), 5), 24); - } - - #[test] - fn edit_low_square_one_output_two_inputs() { - // 2¢ output + 2×2¢ input = 6¢ - assert_eq!(cost_edit(Some(CommonQuality::Low), Some(CommonAspectRatio::Square), Some(1), 2), 6); - } - - #[test] - fn edit_low_tall_four_outputs_one_input() { - // 4×2¢ output + 1×2¢ input = 10¢ - assert_eq!(cost_edit(Some(CommonQuality::Low), Some(CommonAspectRatio::TallNineBySixteen), Some(4), 1), 10); - } - - // ── Default quality in edit mode (None → High) ────────────────────── - - #[test] - fn edit_default_quality_square_one_output_one_input() { - // 17¢ output + 2¢ input = 19¢ - assert_eq!(cost_edit(None, Some(CommonAspectRatio::Square), Some(1), 1), 19); - } - - #[test] - fn edit_default_quality_wide_one_output_one_input() { - // 25¢ output + 2¢ input = 27¢ - assert_eq!(cost_edit(None, Some(CommonAspectRatio::WideSixteenByNine), Some(1), 1), 27); - } - - // ── Quality × size matrix in edit mode (1 output, 1 input) ────────── - - #[test] - fn edit_low_square() { - // 2¢ + 2¢ = 4¢ - assert_eq!(cost_edit(Some(CommonQuality::Low), Some(CommonAspectRatio::Square), Some(1), 1), 4); - } - - #[test] - fn edit_low_wide() { - // 2¢ + 2¢ = 4¢ - assert_eq!(cost_edit(Some(CommonQuality::Low), Some(CommonAspectRatio::WideSixteenByNine), Some(1), 1), 4); - } - - #[test] - fn edit_medium_square() { - // 5¢ + 2¢ = 7¢ - assert_eq!(cost_edit(Some(CommonQuality::Medium), Some(CommonAspectRatio::Square), Some(1), 1), 7); - } - - #[test] - fn edit_medium_wide() { - // 7¢ + 2¢ = 9¢ - assert_eq!(cost_edit(Some(CommonQuality::Medium), Some(CommonAspectRatio::WideSixteenByNine), Some(1), 1), 9); - } - - #[test] - fn edit_medium_tall() { - // 7¢ + 2¢ = 9¢ - assert_eq!(cost_edit(Some(CommonQuality::Medium), Some(CommonAspectRatio::TallNineBySixteen), Some(1), 1), 9); - } - - #[test] - fn edit_high_square() { - // 17¢ + 2¢ = 19¢ - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1), 1), 19); - } - - #[test] - fn edit_high_wide() { - // 25¢ + 2¢ = 27¢ - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::WideSixteenByNine), Some(1), 1), 27); - } - - #[test] - fn edit_high_tall() { - // 25¢ + 2¢ = 27¢ - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::TallNineBySixteen), Some(1), 1), 27); - } - - // ── Batch above max clamps ────────────────────────────────────────── - - #[test] - fn edit_batch_above_max_clamps() { - // High square 4×17 + 1×2 = 70¢ (clamped from 9 to 4) - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(9), 1), 70); - } - - // ── Cost matches text + 2¢ per input image ───────────────────────── - - #[test] - fn edit_and_text_cost_differ_by_input_image_cost() { - for q in [CommonQuality::Low, CommonQuality::Medium, CommonQuality::High] { - let text = distill_text(&make_request(Some("p"), None, Some(q), Some(1), None)) - .cost.cost_in_usd_cents.unwrap(); - let (tokens, hydration) = fake_image_refs(1); - let edit = distill_edit( - &make_request(Some("p"), None, Some(q), Some(1), Some(tokens)), - &hydration, - ).cost.cost_in_usd_cents.unwrap(); - assert_eq!( - edit - text, - 2, - "expected edit to cost 2¢ more than text for quality {:?}", - q, - ); - } - } - - // ── Metadata flags ─────────────────────────────────────────────────── - - #[test] - fn cost_metadata_flags_are_default() { - let (tokens, hydration) = fake_image_refs(1); - let request = make_request(Some("p"), None, None, Some(1), Some(tokens)); - let distilled = distill_edit(&request, &hydration); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_gpt_image_1::{ - FalGptImage1ImageSize, FalGptImage1NumImages, FalGptImage1Quality, PlanFalGptImage1, - }; - - fn with_text_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill_text(request); - match distilled.plan() { - ImageGenerationPlan::FalGptImage1(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalGptImage1, got {:?}", other), - } - } - - fn with_edit_plan( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - assertion: F, - ) { - let distilled = distill_edit(request, hydration); - match distilled.plan() { - ImageGenerationPlan::FalGptImage1(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalGptImage1, got {:?}", other), - } - } - - mod text { - use super::*; - - // ── Quality mapping ─────────────────────────────────────────────────── - - #[test] - fn default_quality_is_high() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage1Quality::High)); - }); - } - - #[test] - fn low_quality_passes_through() { - with_text_plan(&make_request(Some("p"), None, Some(CommonQuality::Low), Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage1Quality::Low)); - }); - } - - #[test] - fn medium_quality_passes_through() { - with_text_plan(&make_request(Some("p"), None, Some(CommonQuality::Medium), Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage1Quality::Medium)); - }); - } - - #[test] - fn high_quality_passes_through() { - with_text_plan(&make_request(Some("p"), None, Some(CommonQuality::High), Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage1Quality::High)); - }); - } - - // ── Mode detection ──────────────────────────────────────────────────── - - #[test] - fn text_mode_has_empty_image_urls() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(plan.image_urls.is_empty()); - }); - } - - // ── Image size mappings ─────────────────────────────────────────────── - - #[test] - fn default_image_size_is_none() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(plan.image_size.is_none()); - }); - } - - #[test] - fn auto_variants_yield_none() { - for ar in [CommonAspectRatio::Auto, CommonAspectRatio::Auto2k, CommonAspectRatio::Auto4k] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(plan.image_size.is_none(), "expected None for {:?}", ar); - }); - } - } - - #[test] - fn square_yields_square() { - for ar in [CommonAspectRatio::Square, CommonAspectRatio::SquareHd] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage1ImageSize::Square)), "expected Square for {:?}", ar); - }); - } - } - - #[test] - fn wide_variants_yield_horizontal() { - let wide_ars = [ - CommonAspectRatio::WideFiveByFour, - CommonAspectRatio::WideFourByThree, - CommonAspectRatio::WideThreeByTwo, - CommonAspectRatio::WideSixteenByNine, - CommonAspectRatio::WideTwentyOneByNine, - CommonAspectRatio::Wide, - ]; - for ar in wide_ars { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage1ImageSize::Horizontal)), "expected Horizontal for {:?}", ar); - }); - } - } - - #[test] - fn tall_variants_yield_vertical() { - let tall_ars = [ - CommonAspectRatio::TallFourByFive, - CommonAspectRatio::TallThreeByFour, - CommonAspectRatio::TallTwoByThree, - CommonAspectRatio::TallNineBySixteen, - CommonAspectRatio::TallNineByTwentyOne, - CommonAspectRatio::Tall, - ]; - for ar in tall_ars { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage1ImageSize::Vertical)), "expected Vertical for {:?}", ar); - }); - } - } - - // ── Num images mapping ──────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - with_text_plan(&make_request(Some("p"), None, None, None, None), |plan| { - assert!(matches!(plan.num_images, FalGptImage1NumImages::One)); - }); - } - - #[test] - fn batch_direct_mapping() { - let cases = [ - (1u16, FalGptImage1NumImages::One), - (2, FalGptImage1NumImages::Two), - (3, FalGptImage1NumImages::Three), - (4, FalGptImage1NumImages::Four), - ]; - for (count, expected) in cases { - with_text_plan(&make_request(Some("p"), None, None, Some(count), None), |plan| { - assert!( - std::mem::discriminant(&plan.num_images) == std::mem::discriminant(&expected), - "expected {:?} for count {}", expected, count, - ); - }); - } - } - - #[test] - fn batch_above_four_clamps_to_four() { - with_text_plan(&make_request(Some("p"), None, None, Some(9), None), |plan| { - assert!(matches!(plan.num_images, FalGptImage1NumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_text_plan(&make_request(Some("a corgi in a hat"), None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_text_plan(&make_request(None, None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } - - mod edit { - use super::*; - - // ── Mode detection / image-url passthrough ──────────────────────────── - - #[test] - fn edit_mode_populates_image_urls() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 1); - assert!(plan.image_urls[0].starts_with("https://fake.example.com/")); - }, - ); - } - - #[test] - fn edit_mode_with_two_image_refs() { - let (tokens, hydration) = fake_image_refs(2); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 2); - }, - ); - } - - #[test] - fn edit_mode_with_five_image_refs() { - let (tokens, hydration) = fake_image_refs(5); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 5); - }, - ); - } - - #[test] - fn edit_image_urls_are_hydrated_from_map() { - let (tokens, hydration) = fake_image_refs(3); - let expected: Vec = tokens.iter() - .map(|t| hydration.get(t).unwrap().to_string()) - .collect(); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens.clone())), - &hydration, - |plan| { - assert_eq!(plan.image_urls, expected); - }, - ); - } - - // ── Quality mapping in edit mode ────────────────────────────────────── - - #[test] - fn default_quality_is_high_in_edit_mode() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert!(matches!(plan.quality, FalGptImage1Quality::High)); - }, - ); - } - - #[test] - fn low_quality_passes_through_in_edit_mode() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(CommonQuality::Low), Some(1), Some(tokens)), - &hydration, - |plan| { - assert!(matches!(plan.quality, FalGptImage1Quality::Low)); - }, - ); - } - - #[test] - fn high_quality_passes_through_in_edit_mode() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(CommonQuality::High), Some(1), Some(tokens)), - &hydration, - |plan| { - assert!(matches!(plan.quality, FalGptImage1Quality::High)); - }, - ); - } - - // ── Image size in edit mode ────────────────────────────────────────── - - #[test] - fn edit_default_image_size_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(plan.image_size.is_none()); }, - ); - } - - #[test] - fn edit_square_yields_square() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Square), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalGptImage1ImageSize::Square))); }, - ); - } - - #[test] - fn edit_wide_yields_horizontal() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::WideSixteenByNine), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalGptImage1ImageSize::Horizontal))); }, - ); - } - - #[test] - fn edit_tall_yields_vertical() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::TallNineBySixteen), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalGptImage1ImageSize::Vertical))); }, - ); - } - - #[test] - fn edit_auto_yields_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Auto), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(plan.image_size.is_none()); }, - ); - } - - // ── Num images in edit mode ────────────────────────────────────────── - - #[test] - fn edit_default_batch_count_is_one() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, None, Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalGptImage1NumImages::One)); }, - ); - } - - #[test] - fn edit_batch_of_four_yields_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(4), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalGptImage1NumImages::Four)); }, - ); - } - - #[test] - fn edit_batch_above_four_clamps_to_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(9), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalGptImage1NumImages::Four)); }, - ); - } - - // ── Prompt passthrough in edit mode ──────────────────────────────────── - - #[test] - fn edit_prompt_is_passed_through() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("make it shiny"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, Some("make it shiny".to_string())); }, - ); - } - - #[test] - fn edit_missing_prompt_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(None, None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, None); }, - ); - } - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1p5.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1p5.rs deleted file mode 100644 index 6b8b255caa..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1p5.rs +++ /dev/null @@ -1,721 +0,0 @@ -//! Tests for omni-gen against the GPT Image 1.5 model. -//! -//! GPT Image 1.5 is multi-function: text-to-image and image-to-image (edit). -//! -//! Artcraft-tier pricing (what the user is billed): -//! Low – 1¢ (square/unset) / 2¢ (wide or tall) -//! Medium – 4¢ (square/unset) / 5¢ (wide) / 6¢ (tall) -//! High – 14¢ (square/unset) / 20¢ (wide or tall) ← default when quality unset -//! -//! The Fal execution plan maps quality from the omni request (defaulting to -//! High when unspecified) and maps aspect ratios to three image sizes: -//! Square (1024×1024), Wide (1536×1024), Tall (1024×1536). -//! Auto/unset maps to None (API default, treated as square for pricing). - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use enums::common::generation::common_quality::CommonQuality; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - quality: Option, - image_batch_count: Option, - image_media_tokens: Option>, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::GptImage1p5), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens, - resolution: None, - aspect_ratio, - quality, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - fn fake_image_refs(n: usize) -> (Vec, HashMap) { - let mut tokens = Vec::with_capacity(n); - let mut map = HashMap::with_capacity(n); - for i in 0..n { - let token = MediaFileToken::new_from_str(&format!("mf_test{:028}", i)); - let url = Url::parse(&format!("https://fake.example.com/img{}.png", i)).unwrap(); - map.insert(token.clone(), url); - tokens.push(token); - } - (tokens, map) - } - - fn distill_text(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for GPT Image 1.5 (text)") - } - - fn distill_edit( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - ) -> DistilledImageRequest { - distill_image_request(request, Some(hydration)) - .expect("distill_image_request should succeed for GPT Image 1.5 (edit)") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST - // ──────────────────────────────────────────────────────────────────────────── - mod cost { - use super::*; - - mod text { - use super::*; - - fn cost( - quality: Option, - aspect_ratio: Option, - batch: Option, - ) -> u64 { - let request = make_request(Some("p"), aspect_ratio, quality, batch, None); - distill_text(&request).cost.cost_in_usd_cents.unwrap() - } - - // ── Default quality (None → High) ──────────────────────────────────── - - #[test] - fn default_quality_square_costs_14() { - assert_eq!(cost(None, Some(CommonAspectRatio::Square), Some(1)), 14); - } - - #[test] - fn default_quality_unset_costs_14() { - assert_eq!(cost(None, None, Some(1)), 14); - } - - #[test] - fn default_quality_auto_costs_14() { - assert_eq!(cost(None, Some(CommonAspectRatio::Auto), Some(1)), 14); - } - - #[test] - fn default_quality_wide_costs_20() { - assert_eq!(cost(None, Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 20); - } - - #[test] - fn default_quality_tall_costs_20() { - assert_eq!(cost(None, Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 20); - } - - // ── Low quality (1¢ square, 2¢ wide/tall) ──────────────────────────── - - #[test] - fn low_square_costs_1() { - assert_eq!(cost(Some(CommonQuality::Low), Some(CommonAspectRatio::Square), Some(1)), 1); - } - - #[test] - fn low_unset_costs_1() { - assert_eq!(cost(Some(CommonQuality::Low), None, Some(1)), 1); - } - - #[test] - fn low_wide_costs_2() { - assert_eq!(cost(Some(CommonQuality::Low), Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 2); - } - - #[test] - fn low_tall_costs_2() { - assert_eq!(cost(Some(CommonQuality::Low), Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 2); - } - - #[test] - fn low_four_images() { - assert_eq!(cost(Some(CommonQuality::Low), None, Some(4)), 4); - } - - #[test] - fn low_batch_above_max_clamps() { - assert_eq!(cost(Some(CommonQuality::Low), None, Some(7)), 4); - } - - // ── Medium quality (4¢ square, 5¢ wide, 6¢ tall) ───────────────────── - - #[test] - fn medium_square_costs_4() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::Square), Some(1)), 4); - } - - #[test] - fn medium_unset_costs_4() { - assert_eq!(cost(Some(CommonQuality::Medium), None, Some(1)), 4); - } - - #[test] - fn medium_wide_costs_5() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 5); - } - - #[test] - fn medium_tall_costs_6() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 6); - } - - #[test] - fn medium_square_four_images() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::Square), Some(4)), 16); - } - - #[test] - fn medium_wide_four_images() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::WideSixteenByNine), Some(4)), 20); - } - - #[test] - fn medium_tall_four_images() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::TallNineBySixteen), Some(4)), 24); - } - - // ── High quality (14¢ square, 20¢ wide/tall) ───────────────────────── - - #[test] - fn high_square_costs_14() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1)), 14); - } - - #[test] - fn high_unset_costs_14() { - assert_eq!(cost(Some(CommonQuality::High), None, Some(1)), 14); - } - - #[test] - fn high_wide_costs_20() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 20); - } - - #[test] - fn high_tall_costs_20() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 20); - } - - #[test] - fn high_square_four_images() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(4)), 56); - } - - #[test] - fn high_wide_four_images() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::WideSixteenByNine), Some(4)), 80); - } - - #[test] - fn high_tall_four_images() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::TallNineBySixteen), Some(4)), 80); - } - - // ── Batch scaling ──────────────────────────────────────────────────── - - #[test] - fn high_square_batch_scaling() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1)), 14); - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(2)), 28); - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(3)), 42); - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(4)), 56); - } - - // ── Metadata flags ─────────────────────────────────────────────────── - - #[test] - fn cost_metadata_flags_are_default() { - let distilled = distill_text(&make_request(Some("p"), None, None, Some(1), None)); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - - #[test] - fn cost_is_independent_of_prompt() { - let a = distill_text(&make_request(Some("a cat"), None, None, Some(1), None)); - let b = distill_text(&make_request(None, None, None, Some(1), None)); - assert_eq!(a.cost.cost_in_usd_cents, b.cost.cost_in_usd_cents); - } - } - - mod edit { - use super::*; - - fn cost_edit( - quality: Option, - aspect_ratio: Option, - batch: Option, - num_input_images: usize, - ) -> u64 { - let (tokens, hydration) = fake_image_refs(num_input_images); - let request = make_request(Some("p"), aspect_ratio, quality, batch, Some(tokens)); - distill_edit(&request, &hydration).cost.cost_in_usd_cents.unwrap() - } - - // ── Full quality × size matrix (1 output, 1 input) ───────────────── - - #[test] - fn edit_low_square() { assert_eq!(cost_edit(Some(CommonQuality::Low), Some(CommonAspectRatio::Square), Some(1), 1), 1); } - - #[test] - fn edit_low_wide() { assert_eq!(cost_edit(Some(CommonQuality::Low), Some(CommonAspectRatio::WideSixteenByNine), Some(1), 1), 2); } - - #[test] - fn edit_low_tall() { assert_eq!(cost_edit(Some(CommonQuality::Low), Some(CommonAspectRatio::TallNineBySixteen), Some(1), 1), 2); } - - #[test] - fn edit_medium_square() { assert_eq!(cost_edit(Some(CommonQuality::Medium), Some(CommonAspectRatio::Square), Some(1), 1), 4); } - - #[test] - fn edit_medium_wide() { assert_eq!(cost_edit(Some(CommonQuality::Medium), Some(CommonAspectRatio::WideSixteenByNine), Some(1), 1), 5); } - - #[test] - fn edit_medium_tall() { assert_eq!(cost_edit(Some(CommonQuality::Medium), Some(CommonAspectRatio::TallNineBySixteen), Some(1), 1), 6); } - - #[test] - fn edit_high_square() { assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1), 1), 14); } - - #[test] - fn edit_high_wide() { assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::WideSixteenByNine), Some(1), 1), 20); } - - #[test] - fn edit_high_tall() { assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::TallNineBySixteen), Some(1), 1), 20); } - - // ── Default quality in edit mode ──────────────────────────────────── - - #[test] - fn edit_default_quality_square() { assert_eq!(cost_edit(None, Some(CommonAspectRatio::Square), Some(1), 1), 14); } - - #[test] - fn edit_default_quality_wide() { assert_eq!(cost_edit(None, Some(CommonAspectRatio::WideSixteenByNine), Some(1), 1), 20); } - - // ── Cost is independent of input image count (no per-input surcharge) - - #[test] - fn edit_cost_independent_of_input_image_count() { - for num_refs in [1usize, 2, 3, 5] { - assert_eq!( - cost_edit(Some(CommonQuality::Medium), None, Some(2), num_refs), - 8, - "expected 8¢ regardless of {} input images", - num_refs, - ); - } - } - - // ── Batch scaling in edit mode ────────────────────────────────────── - - #[test] - fn edit_high_square_four_outputs() { - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(4), 1), 56); - } - - #[test] - fn edit_batch_above_max_clamps() { - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(9), 1), 56); - } - - // ── Edit and text costs match ─────────────────────────────────────── - - #[test] - fn edit_and_text_cost_match() { - for q in [CommonQuality::Low, CommonQuality::Medium, CommonQuality::High] { - let text = distill_text(&make_request(Some("p"), None, Some(q), Some(2), None)) - .cost.cost_in_usd_cents.unwrap(); - let (tokens, hydration) = fake_image_refs(3); - let edit = distill_edit( - &make_request(Some("p"), None, Some(q), Some(2), Some(tokens)), - &hydration, - ).cost.cost_in_usd_cents.unwrap(); - assert_eq!(text, edit, "text/edit cost should match for quality {:?}", q); - } - } - - // ── Metadata flags ────────────────────────────────────────────────── - - #[test] - fn cost_metadata_flags_are_default() { - let (tokens, hydration) = fake_image_refs(1); - let distilled = distill_edit(&make_request(Some("p"), None, None, Some(1), Some(tokens)), &hydration); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_gpt_image_1p5::{ - FalGptImage1p5ImageSize, FalGptImage1p5NumImages, FalGptImage1p5Quality, PlanFalGptImage1p5, - }; - - fn with_text_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill_text(request); - match distilled.plan() { - ImageGenerationPlan::FalGptImage1p5(plan) => assertion(plan), - other => panic!("expected FalGptImage1p5, got {:?}", other), - } - } - - fn with_edit_plan( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - assertion: F, - ) { - let distilled = distill_edit(request, hydration); - match distilled.plan() { - ImageGenerationPlan::FalGptImage1p5(plan) => assertion(plan), - other => panic!("expected FalGptImage1p5, got {:?}", other), - } - } - - mod text { - use super::*; - - // ── Quality mapping ─────────────────────────────────────────────────── - - #[test] - fn default_quality_is_high() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage1p5Quality::High)); - }); - } - - #[test] - fn low_quality_passes_through() { - with_text_plan(&make_request(Some("p"), None, Some(CommonQuality::Low), Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage1p5Quality::Low)); - }); - } - - #[test] - fn medium_quality_passes_through() { - with_text_plan(&make_request(Some("p"), None, Some(CommonQuality::Medium), Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage1p5Quality::Medium)); - }); - } - - #[test] - fn high_quality_passes_through() { - with_text_plan(&make_request(Some("p"), None, Some(CommonQuality::High), Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage1p5Quality::High)); - }); - } - - // ── Mode detection ──────────────────────────────────────────────────── - - #[test] - fn text_mode_has_empty_image_urls() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(plan.image_urls.is_empty()); - }); - } - - // ── Image size mappings ─────────────────────────────────────────────── - - #[test] - fn default_image_size_is_none() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(plan.image_size.is_none()); - }); - } - - #[test] - fn auto_variants_yield_none() { - for ar in [CommonAspectRatio::Auto, CommonAspectRatio::Auto2k, CommonAspectRatio::Auto4k] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(plan.image_size.is_none(), "expected None for {:?}", ar); - }); - } - } - - #[test] - fn square_yields_square() { - for ar in [CommonAspectRatio::Square, CommonAspectRatio::SquareHd] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage1p5ImageSize::Square))); - }); - } - } - - #[test] - fn wide_variants_yield_wide() { - for ar in [ - CommonAspectRatio::WideFiveByFour, CommonAspectRatio::WideFourByThree, - CommonAspectRatio::WideThreeByTwo, CommonAspectRatio::WideSixteenByNine, - CommonAspectRatio::WideTwentyOneByNine, CommonAspectRatio::Wide, - ] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage1p5ImageSize::Wide)), "expected Wide for {:?}", ar); - }); - } - } - - #[test] - fn tall_variants_yield_tall() { - for ar in [ - CommonAspectRatio::TallFourByFive, CommonAspectRatio::TallThreeByFour, - CommonAspectRatio::TallTwoByThree, CommonAspectRatio::TallNineBySixteen, - CommonAspectRatio::TallNineByTwentyOne, CommonAspectRatio::Tall, - ] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage1p5ImageSize::Tall)), "expected Tall for {:?}", ar); - }); - } - } - - // ── Num images mapping ──────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - with_text_plan(&make_request(Some("p"), None, None, None, None), |plan| { - assert!(matches!(plan.num_images, FalGptImage1p5NumImages::One)); - }); - } - - #[test] - fn batch_direct_mapping() { - for (count, expected) in [ - (1u16, FalGptImage1p5NumImages::One), (2, FalGptImage1p5NumImages::Two), - (3, FalGptImage1p5NumImages::Three), (4, FalGptImage1p5NumImages::Four), - ] { - with_text_plan(&make_request(Some("p"), None, None, Some(count), None), |plan| { - assert!(std::mem::discriminant(&plan.num_images) == std::mem::discriminant(&expected)); - }); - } - } - - #[test] - fn batch_above_four_clamps_to_four() { - with_text_plan(&make_request(Some("p"), None, None, Some(9), None), |plan| { - assert!(matches!(plan.num_images, FalGptImage1p5NumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_text_plan(&make_request(Some("a corgi in a hat"), None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_text_plan(&make_request(None, None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } - - mod edit { - use super::*; - - // ── Mode detection / image-url passthrough ──────────────────────────── - - #[test] - fn edit_mode_populates_image_urls() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 1); - assert!(plan.image_urls[0].starts_with("https://fake.example.com/")); - }, - ); - } - - #[test] - fn edit_mode_with_three_image_refs() { - let (tokens, hydration) = fake_image_refs(3); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.image_urls.len(), 3); }, - ); - } - - #[test] - fn edit_mode_with_five_image_refs() { - let (tokens, hydration) = fake_image_refs(5); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.image_urls.len(), 5); }, - ); - } - - #[test] - fn edit_image_urls_are_hydrated_from_map() { - let (tokens, hydration) = fake_image_refs(3); - let expected: Vec = tokens.iter() - .map(|t| hydration.get(t).unwrap().to_string()) - .collect(); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens.clone())), - &hydration, - |plan| { assert_eq!(plan.image_urls, expected); }, - ); - } - - // ── Quality mapping in edit mode ────────────────────────────────────── - - #[test] - fn default_quality_is_high_in_edit_mode() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.quality, FalGptImage1p5Quality::High)); }, - ); - } - - #[test] - fn low_quality_passes_through_in_edit_mode() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(CommonQuality::Low), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.quality, FalGptImage1p5Quality::Low)); }, - ); - } - - #[test] - fn high_quality_passes_through_in_edit_mode() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(CommonQuality::High), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.quality, FalGptImage1p5Quality::High)); }, - ); - } - - // ── Image size in edit mode ────────────────────────────────────────── - - #[test] - fn edit_default_image_size_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(plan.image_size.is_none()); }, - ); - } - - #[test] - fn edit_square_yields_square() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Square), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalGptImage1p5ImageSize::Square))); }, - ); - } - - #[test] - fn edit_wide_yields_wide() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::WideSixteenByNine), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalGptImage1p5ImageSize::Wide))); }, - ); - } - - #[test] - fn edit_tall_yields_tall() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::TallNineBySixteen), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalGptImage1p5ImageSize::Tall))); }, - ); - } - - // ── Num images in edit mode ────────────────────────────────────────── - - #[test] - fn edit_default_batch_count_is_one() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, None, Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalGptImage1p5NumImages::One)); }, - ); - } - - #[test] - fn edit_batch_of_four_yields_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(4), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalGptImage1p5NumImages::Four)); }, - ); - } - - #[test] - fn edit_batch_above_four_clamps_to_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(9), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalGptImage1p5NumImages::Four)); }, - ); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn edit_prompt_is_passed_through() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("make it shiny"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, Some("make it shiny".to_string())); }, - ); - } - - #[test] - fn edit_missing_prompt_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(None, None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, None); }, - ); - } - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_2.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_2.rs deleted file mode 100644 index 68dff3f116..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_2.rs +++ /dev/null @@ -1,655 +0,0 @@ -//! Tests for omni-gen against the GPT Image 2 model. -//! -//! GPT Image 2 is multi-function: text-to-image and image-to-image (edit). -//! -//! Artcraft-tier pricing (what the user is billed): -//! Low – 1¢ for all sizes -//! Medium – 4¢ (landscape_4_3/portrait_4_3/landscape_16_9/portrait_16_9) -//! 6¢ (square/square_hd/auto/unset) -//! High – 15¢ (landscape_4_3/portrait_4_3) -//! 16¢ (landscape_16_9/portrait_16_9) -//! 22¢ (square/unset) -//! 23¢ (square_hd/auto) ← default when quality unset -//! -//! The Fal execution plan maps quality from the omni request (defaulting to -//! High when unspecified) and maps aspect ratios to six image sizes plus Auto: -//! SquareHd, Square, Portrait4x3, Portrait16x9, Landscape4x3, Landscape16x9. -//! Auto/unset maps to Auto (treated as square_hd for pricing). - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use enums::common::generation::common_quality::CommonQuality; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ──────────────────────────────────────────────────────────────────────────── - // COST - // ──────────────────────────────────────────────────────────────────────────── - mod cost { - use super::*; - - mod text { - use super::*; - - fn cost( - quality: Option, - aspect_ratio: Option, - batch: Option, - ) -> u64 { - let request = make_request(Some("p"), aspect_ratio, quality, batch, None); - distill_text(&request).cost.cost_in_usd_cents.unwrap() - } - - // ── Default quality (None → High) ── - - #[test] - fn default_quality_square_costs_22() { - assert_eq!(cost(None, Some(CommonAspectRatio::Square), Some(1)), 22); - } - - #[test] - fn default_quality_unset_costs_23() { - // unset → Auto → 23¢ - assert_eq!(cost(None, None, Some(1)), 23); - } - - #[test] - fn default_quality_auto_costs_23() { - assert_eq!(cost(None, Some(CommonAspectRatio::Auto), Some(1)), 23); - } - - #[test] - fn default_quality_square_hd_costs_23() { - assert_eq!(cost(None, Some(CommonAspectRatio::SquareHd), Some(1)), 23); - } - - #[test] - fn default_quality_wide_4x3_costs_15() { - assert_eq!(cost(None, Some(CommonAspectRatio::WideFourByThree), Some(1)), 15); - } - - #[test] - fn default_quality_wide_16x9_costs_16() { - assert_eq!(cost(None, Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 16); - } - - #[test] - fn default_quality_tall_4x5_costs_15() { - assert_eq!(cost(None, Some(CommonAspectRatio::TallFourByFive), Some(1)), 15); - } - - #[test] - fn default_quality_tall_16x9_costs_16() { - assert_eq!(cost(None, Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 16); - } - - // ── Low quality (1¢ all sizes) ── - - #[test] - fn low_square_costs_1() { - assert_eq!(cost(Some(CommonQuality::Low), Some(CommonAspectRatio::Square), Some(1)), 1); - } - - #[test] - fn low_unset_costs_1() { - assert_eq!(cost(Some(CommonQuality::Low), None, Some(1)), 1); - } - - #[test] - fn low_wide_costs_1() { - assert_eq!(cost(Some(CommonQuality::Low), Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 1); - } - - #[test] - fn low_tall_costs_1() { - assert_eq!(cost(Some(CommonQuality::Low), Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 1); - } - - #[test] - fn low_four_images() { - assert_eq!(cost(Some(CommonQuality::Low), None, Some(4)), 4); - } - - #[test] - fn low_batch_above_max_clamps() { - assert_eq!(cost(Some(CommonQuality::Low), None, Some(7)), 4); - } - - // ── Medium quality (4¢ 4:3/16:9, 6¢ square/square_hd/auto) ── - - #[test] - fn medium_square_costs_6() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::Square), Some(1)), 6); - } - - #[test] - fn medium_square_hd_costs_6() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::SquareHd), Some(1)), 6); - } - - #[test] - fn medium_unset_costs_6() { - assert_eq!(cost(Some(CommonQuality::Medium), None, Some(1)), 6); - } - - #[test] - fn medium_wide_4x3_costs_4() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::WideFourByThree), Some(1)), 4); - } - - #[test] - fn medium_wide_16x9_costs_4() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 4); - } - - #[test] - fn medium_tall_4x3_costs_4() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::TallThreeByFour), Some(1)), 4); - } - - #[test] - fn medium_tall_16x9_costs_4() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 4); - } - - #[test] - fn medium_square_four_images() { - assert_eq!(cost(Some(CommonQuality::Medium), Some(CommonAspectRatio::Square), Some(4)), 24); - } - - // ── High quality (15¢ 4:3, 16¢ 16:9, 22¢ square, 23¢ square_hd/auto) ── - - #[test] - fn high_square_costs_22() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1)), 22); - } - - #[test] - fn high_square_hd_costs_23() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::SquareHd), Some(1)), 23); - } - - #[test] - fn high_unset_costs_23() { - assert_eq!(cost(Some(CommonQuality::High), None, Some(1)), 23); - } - - #[test] - fn high_wide_4x3_costs_15() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::WideFourByThree), Some(1)), 15); - } - - #[test] - fn high_wide_16x9_costs_16() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::WideSixteenByNine), Some(1)), 16); - } - - #[test] - fn high_tall_4x3_costs_15() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::TallThreeByFour), Some(1)), 15); - } - - #[test] - fn high_tall_16x9_costs_16() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::TallNineBySixteen), Some(1)), 16); - } - - #[test] - fn high_square_four_images() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(4)), 88); - } - - #[test] - fn high_square_hd_four_images() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::SquareHd), Some(4)), 92); - } - - // ── Batch scaling ── - - #[test] - fn high_square_batch_scaling() { - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1)), 22); - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(2)), 44); - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(3)), 66); - assert_eq!(cost(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(4)), 88); - } - - // ── Metadata flags ── - - #[test] - fn cost_metadata_flags_are_default() { - let distilled = distill_text(&make_request(Some("p"), None, None, Some(1), None)); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - - #[test] - fn cost_is_independent_of_prompt() { - let a = distill_text(&make_request(Some("a cat"), None, None, Some(1), None)); - let b = distill_text(&make_request(None, None, None, Some(1), None)); - assert_eq!(a.cost.cost_in_usd_cents, b.cost.cost_in_usd_cents); - } - } - - mod edit { - use super::*; - - fn cost_edit( - quality: Option, - aspect_ratio: Option, - batch: Option, - num_input_images: usize, - ) -> u64 { - let (tokens, hydration) = fake_image_refs(num_input_images); - let request = make_request(Some("p"), aspect_ratio, quality, batch, Some(tokens)); - distill_edit(&request, &hydration).cost.cost_in_usd_cents.unwrap() - } - - #[test] - fn edit_low_square() { assert_eq!(cost_edit(Some(CommonQuality::Low), Some(CommonAspectRatio::Square), Some(1), 1), 1); } - - #[test] - fn edit_medium_square() { assert_eq!(cost_edit(Some(CommonQuality::Medium), Some(CommonAspectRatio::Square), Some(1), 1), 6); } - - #[test] - fn edit_medium_wide_4x3() { assert_eq!(cost_edit(Some(CommonQuality::Medium), Some(CommonAspectRatio::WideFourByThree), Some(1), 1), 4); } - - #[test] - fn edit_high_square() { assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(1), 1), 22); } - - #[test] - fn edit_high_square_hd() { assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::SquareHd), Some(1), 1), 23); } - - #[test] - fn edit_high_wide_16x9() { assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::WideSixteenByNine), Some(1), 1), 16); } - - #[test] - fn edit_high_tall_4x3() { assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::TallThreeByFour), Some(1), 1), 15); } - - #[test] - fn edit_default_quality_square() { assert_eq!(cost_edit(None, Some(CommonAspectRatio::Square), Some(1), 1), 22); } - - #[test] - fn edit_cost_independent_of_input_image_count() { - for num_refs in [1usize, 2, 3, 5] { - assert_eq!( - cost_edit(Some(CommonQuality::Medium), None, Some(2), num_refs), - 12, - "expected 12¢ regardless of {} input images", - num_refs, - ); - } - } - - #[test] - fn edit_high_square_four_outputs() { - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(4), 1), 88); - } - - #[test] - fn edit_batch_above_max_clamps() { - assert_eq!(cost_edit(Some(CommonQuality::High), Some(CommonAspectRatio::Square), Some(9), 1), 88); - } - - #[test] - fn edit_and_text_cost_match() { - for q in [CommonQuality::Low, CommonQuality::Medium, CommonQuality::High] { - let text = distill_text(&make_request(Some("p"), None, Some(q), Some(2), None)) - .cost.cost_in_usd_cents.unwrap(); - let (tokens, hydration) = fake_image_refs(3); - let edit = distill_edit( - &make_request(Some("p"), None, Some(q), Some(2), Some(tokens)), - &hydration, - ).cost.cost_in_usd_cents.unwrap(); - assert_eq!(text, edit, "text/edit cost should match for quality {:?}", q); - } - } - - #[test] - fn cost_metadata_flags_are_default() { - let (tokens, hydration) = fake_image_refs(1); - let distilled = distill_edit(&make_request(Some("p"), None, None, Some(1), Some(tokens)), &hydration); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_gpt_image_2::{ - FalGptImage2ImageSize, FalGptImage2NumImages, FalGptImage2Quality, PlanFalGptImage2, - }; - - mod text { - use super::*; - - // ── Quality mapping ── - - #[test] - fn default_quality_is_high() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage2Quality::High)); - }); - } - - #[test] - fn low_quality_passes_through() { - with_text_plan(&make_request(Some("p"), None, Some(CommonQuality::Low), Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage2Quality::Low)); - }); - } - - #[test] - fn medium_quality_passes_through() { - with_text_plan(&make_request(Some("p"), None, Some(CommonQuality::Medium), Some(1), None), |plan| { - assert!(matches!(plan.quality, FalGptImage2Quality::Medium)); - }); - } - - // ── Mode detection ── - - #[test] - fn text_mode_has_empty_image_urls() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(plan.image_urls.is_empty()); - }); - } - - // ── Image size mappings ── - - #[test] - fn default_image_size_is_auto() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage2ImageSize::Auto))); - }); - } - - #[test] - fn auto_variants_yield_auto() { - for ar in [CommonAspectRatio::Auto, CommonAspectRatio::Auto2k, CommonAspectRatio::Auto4k] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage2ImageSize::Auto)), "expected Auto for {:?}", ar); - }); - } - } - - #[test] - fn square_yields_square() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Square), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage2ImageSize::Square))); - }); - } - - #[test] - fn square_hd_yields_square_hd() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::SquareHd), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage2ImageSize::SquareHd))); - }); - } - - #[test] - fn wide_4x3_yields_landscape_4x3() { - for ar in [CommonAspectRatio::WideFourByThree, CommonAspectRatio::WideFiveByFour] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage2ImageSize::Landscape4x3)), "expected Landscape4x3 for {:?}", ar); - }); - } - } - - #[test] - fn wide_16x9_yields_landscape_16x9() { - for ar in [CommonAspectRatio::WideSixteenByNine, CommonAspectRatio::Wide, CommonAspectRatio::WideThreeByTwo] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage2ImageSize::Landscape16x9)), "expected Landscape16x9 for {:?}", ar); - }); - } - } - - #[test] - fn tall_4x3_yields_portrait_4x3() { - for ar in [CommonAspectRatio::TallThreeByFour, CommonAspectRatio::TallFourByFive] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage2ImageSize::Portrait4x3)), "expected Portrait4x3 for {:?}", ar); - }); - } - } - - #[test] - fn tall_16x9_yields_portrait_16x9() { - for ar in [CommonAspectRatio::TallNineBySixteen, CommonAspectRatio::Tall, CommonAspectRatio::TallTwoByThree] { - with_text_plan(&make_request(Some("p"), Some(ar), None, Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalGptImage2ImageSize::Portrait16x9)), "expected Portrait16x9 for {:?}", ar); - }); - } - } - - // ── Num images mapping ── - - #[test] - fn default_batch_count_is_one() { - with_text_plan(&make_request(Some("p"), None, None, None, None), |plan| { - assert!(matches!(plan.num_images, FalGptImage2NumImages::One)); - }); - } - - #[test] - fn batch_direct_mapping() { - for (count, expected) in [ - (1u16, FalGptImage2NumImages::One), (2, FalGptImage2NumImages::Two), - (3, FalGptImage2NumImages::Three), (4, FalGptImage2NumImages::Four), - ] { - with_text_plan(&make_request(Some("p"), None, None, Some(count), None), |plan| { - assert!(std::mem::discriminant(&plan.num_images) == std::mem::discriminant(&expected)); - }); - } - } - - #[test] - fn batch_above_four_clamps_to_four() { - with_text_plan(&make_request(Some("p"), None, None, Some(9), None), |plan| { - assert!(matches!(plan.num_images, FalGptImage2NumImages::Four)); - }); - } - - // ── Prompt passthrough ── - - #[test] - fn prompt_is_passed_through() { - with_text_plan(&make_request(Some("a corgi in a hat"), None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_text_plan(&make_request(None, None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } - - mod edit { - use super::*; - - #[test] - fn edit_mode_populates_image_urls() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 1); - assert!(plan.image_urls[0].starts_with("https://fake.example.com/")); - }, - ); - } - - #[test] - fn edit_mode_with_three_image_refs() { - let (tokens, hydration) = fake_image_refs(3); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.image_urls.len(), 3); }, - ); - } - - #[test] - fn edit_image_urls_are_hydrated_from_map() { - let (tokens, hydration) = fake_image_refs(3); - let expected: Vec = tokens.iter() - .map(|t| hydration.get(t).unwrap().to_string()) - .collect(); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens.clone())), - &hydration, - |plan| { assert_eq!(plan.image_urls, expected); }, - ); - } - - #[test] - fn default_quality_is_high_in_edit_mode() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.quality, FalGptImage2Quality::High)); }, - ); - } - - #[test] - fn edit_square_yields_square() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Square), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalGptImage2ImageSize::Square))); }, - ); - } - - #[test] - fn edit_wide_16x9_yields_landscape_16x9() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::WideSixteenByNine), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalGptImage2ImageSize::Landscape16x9))); }, - ); - } - - #[test] - fn edit_batch_above_four_clamps_to_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(9), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalGptImage2NumImages::Four)); }, - ); - } - - #[test] - fn edit_prompt_is_passed_through() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("make it shiny"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, Some("make it shiny".to_string())); }, - ); - } - } - - // ── Plan helpers ── - - fn with_text_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill_text(request); - match distilled.plan() { - ImageGenerationPlan::FalGptImage2(plan) => assertion(plan), - other => panic!("expected FalGptImage2, got {:?}", other), - } - } - - fn with_edit_plan( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - assertion: F, - ) { - let distilled = distill_edit(request, hydration); - match distilled.plan() { - ImageGenerationPlan::FalGptImage2(plan) => assertion(plan), - other => panic!("expected FalGptImage2, got {:?}", other), - } - } - } - - // ── Helpers ── - - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - quality: Option, - image_batch_count: Option, - image_media_tokens: Option>, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::GptImage2), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens, - resolution: None, - aspect_ratio, - quality, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - fn fake_image_refs(n: usize) -> (Vec, HashMap) { - let mut tokens = Vec::with_capacity(n); - let mut map = HashMap::with_capacity(n); - for i in 0..n { - let token = MediaFileToken::new_from_str(&format!("mf_test{:028}", i)); - let url = Url::parse(&format!("https://fake.example.com/img{}.png", i)).unwrap(); - map.insert(token.clone(), url); - tokens.push(token); - } - (tokens, map) - } - - fn distill_text(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for GPT Image 2 (text)") - } - - fn distill_edit( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - ) -> DistilledImageRequest { - distill_image_request(request, Some(hydration)) - .expect("distill_image_request should succeed for GPT Image 2 (edit)") - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana.rs deleted file mode 100644 index fef4177b39..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana.rs +++ /dev/null @@ -1,948 +0,0 @@ -//! Tests for omni-gen against the Nano Banana model. -//! -//! Nano Banana is a multi-function model: it does both text-to-image and -//! image-to-image (edit). The omni-gen distillation routes the same model -//! through two different Fal endpoints depending on whether image refs are -//! supplied. We exercise both modes here. -//! -//! Pricing matches the legacy `nano_banana_multi_function_image_gen_handler`: -//! $0.04 / image (4 credits per image; 1 credit == 1 USD cent) -//! for both text-to-image and image-edit modes — same as -//! `Gemini25FlashTextToImageArgs::calculate_cost_in_cents` and -//! `Gemini25FlashEditArgs::calculate_cost_in_cents`. - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - /// Build a minimal Nano Banana omni request. - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - image_batch_count: Option, - image_media_tokens: Option>, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::NanoBanana), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens, - resolution: None, - aspect_ratio, - quality: None, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - /// Make `n` fake media tokens paired with fake URLs, returned as both a - /// vec (for `image_media_tokens`) and a hydration map. - fn fake_image_refs(n: usize) -> (Vec, HashMap) { - let mut tokens = Vec::with_capacity(n); - let mut map = HashMap::with_capacity(n); - for i in 0..n { - let token = MediaFileToken::new_from_str(&format!("mf_test{:028}", i)); - let url = Url::parse(&format!("https://fake.example.com/img{}.png", i)).unwrap(); - map.insert(token.clone(), url); - tokens.push(token); - } - (tokens, map) - } - - /// Distill a text-to-image request (no image refs, empty hydration map). - fn distill_text(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for Nano Banana (text)") - } - - /// Distill an edit request, providing the hydration map. - fn distill_edit( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - ) -> DistilledImageRequest { - distill_image_request(request, Some(hydration)) - .expect("distill_image_request should succeed for Nano Banana (edit)") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST - // ──────────────────────────────────────────────────────────────────────────── - // - // Pricing comes from `estimate_image_cost_artcraft_nano_banana`: - // $0.04 / image (4 credits per image), independent of mode. - // - // Matches the legacy handler's call to `calculate_cost_in_cents` on both - // Gemini25FlashTextToImageArgs and Gemini25FlashEditArgs. - mod cost { - use super::*; - - mod text { - use super::*; - - fn cost_for_batch(image_batch_count: Option) -> (Option, Option) { - let request = make_request(Some("a happy puppy"), None, image_batch_count, None); - let distilled = distill_text(&request); - ( - distilled.cost.cost_in_credits, - distilled.cost.cost_in_usd_cents, - ) - } - - #[test] - fn default_batch_costs_4_cents() { - let (credits, cents) = cost_for_batch(None); - assert_eq!(credits, Some(4)); - assert_eq!(cents, Some(4)); - } - - #[test] - fn batch_of_one_costs_4_cents() { - let (credits, cents) = cost_for_batch(Some(1)); - assert_eq!(credits, Some(4)); - assert_eq!(cents, Some(4)); - } - - #[test] - fn batch_of_two_costs_8_cents() { - let (credits, cents) = cost_for_batch(Some(2)); - assert_eq!(credits, Some(8)); - assert_eq!(cents, Some(8)); - } - - #[test] - fn batch_of_three_costs_12_cents() { - let (credits, cents) = cost_for_batch(Some(3)); - assert_eq!(credits, Some(12)); - assert_eq!(cents, Some(12)); - } - - #[test] - fn batch_of_four_costs_16_cents() { - let (credits, cents) = cost_for_batch(Some(4)); - assert_eq!(credits, Some(16)); - assert_eq!(cents, Some(16)); - } - - #[test] - fn batch_above_max_clamps_and_costs_16_cents() { - // PayMoreUpgrade clamps over-large batches to the max of 4. - let (credits, cents) = cost_for_batch(Some(7)); - assert_eq!(credits, Some(16)); - assert_eq!(cents, Some(16)); - } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let aspect_ratios = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::SquareHd), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::TallNineBySixteen), - Some(CommonAspectRatio::Auto), - ]; - for ar in aspect_ratios { - let request = make_request(Some("p"), ar, Some(2), None); - let distilled = distill_text(&request); - assert_eq!( - distilled.cost.cost_in_usd_cents, - Some(8), - "expected 8¢ regardless of aspect ratio (got {:?})", - ar, - ); - } - } - - #[test] - fn cost_is_independent_of_prompt() { - let with_prompt = distill_text(&make_request(Some("a cat"), None, Some(3), None)); - let without_prompt = distill_text(&make_request(None, None, Some(3), None)); - assert_eq!(with_prompt.cost.cost_in_usd_cents, Some(12)); - assert_eq!(without_prompt.cost.cost_in_usd_cents, Some(12)); - } - - #[test] - fn cost_metadata_flags_are_default() { - let request = make_request(Some("p"), None, Some(1), None); - let distilled = distill_text(&request); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - } - - mod edit { - use super::*; - - fn cost_for_batch_with_image_refs( - image_batch_count: Option, - num_image_refs: usize, - ) -> (Option, Option) { - let (tokens, hydration) = fake_image_refs(num_image_refs); - let request = make_request(Some("make it pop"), None, image_batch_count, Some(tokens)); - let distilled = distill_edit(&request, &hydration); - ( - distilled.cost.cost_in_credits, - distilled.cost.cost_in_usd_cents, - ) - } - - #[test] - fn default_batch_costs_4_cents() { - let (credits, cents) = cost_for_batch_with_image_refs(None, 1); - assert_eq!(credits, Some(4)); - assert_eq!(cents, Some(4)); - } - - #[test] - fn batch_of_one_costs_4_cents() { - let (credits, cents) = cost_for_batch_with_image_refs(Some(1), 1); - assert_eq!(credits, Some(4)); - assert_eq!(cents, Some(4)); - } - - #[test] - fn batch_of_two_costs_8_cents() { - let (credits, cents) = cost_for_batch_with_image_refs(Some(2), 1); - assert_eq!(credits, Some(8)); - assert_eq!(cents, Some(8)); - } - - #[test] - fn batch_of_three_costs_12_cents() { - let (credits, cents) = cost_for_batch_with_image_refs(Some(3), 1); - assert_eq!(credits, Some(12)); - assert_eq!(cents, Some(12)); - } - - #[test] - fn batch_of_four_costs_16_cents() { - let (credits, cents) = cost_for_batch_with_image_refs(Some(4), 1); - assert_eq!(credits, Some(16)); - assert_eq!(cents, Some(16)); - } - - #[test] - fn batch_above_max_clamps_and_costs_16_cents() { - let (credits, cents) = cost_for_batch_with_image_refs(Some(9), 1); - assert_eq!(credits, Some(16)); - assert_eq!(cents, Some(16)); - } - - #[test] - fn cost_is_independent_of_image_ref_count() { - // Whether the user attaches 1, 2, or 5 reference images, the price is - // the same: $0.04 × num_output_images. - for num_refs in [1usize, 2, 3, 5] { - let (_credits, cents) = cost_for_batch_with_image_refs(Some(2), num_refs); - assert_eq!( - cents, - Some(8), - "expected 8¢ regardless of {} image refs", - num_refs, - ); - } - } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let aspect_ratios = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::Auto), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::TallNineBySixteen), - ]; - for ar in aspect_ratios { - let (tokens, hydration) = fake_image_refs(2); - let request = make_request(Some("p"), ar, Some(2), Some(tokens)); - let distilled = distill_edit(&request, &hydration); - assert_eq!( - distilled.cost.cost_in_usd_cents, - Some(8), - "expected 8¢ regardless of aspect ratio (got {:?})", - ar, - ); - } - } - - #[test] - fn cost_metadata_flags_are_default() { - let (tokens, hydration) = fake_image_refs(1); - let request = make_request(Some("p"), None, Some(1), Some(tokens)); - let distilled = distill_edit(&request, &hydration); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - - #[test] - fn edit_and_text_cost_match_for_same_batch() { - // Both modes are billed at the same $0.04/image rate. - for batch in [1u16, 2, 3, 4] { - let text = distill_text(&make_request(Some("p"), None, Some(batch), None)); - let (tokens, hydration) = fake_image_refs(1); - let edit = distill_edit( - &make_request(Some("p"), None, Some(batch), Some(tokens)), - &hydration, - ); - assert_eq!( - text.cost.cost_in_usd_cents, - edit.cost.cost_in_usd_cents, - "text/edit cost diverged at batch={}", - batch, - ); - } - } - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - // - // The distilled plan is `ImageGenerationPlan::FalNanoBanana(PlanFalNanoBanana)`. - // This single plan struct carries both: - // - `t2i_aspect_ratio: Option` - // - `edit_aspect_ratio: Option` - // - `image_urls: Vec` — empty in text mode, populated in edit mode - // - `num_images: FalNbNumImages` - // - // The handler picks t2i vs edit based on whether `image_urls` is empty. - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana::{ - FalNbNumImages, PlanFalNanoBanana, - }; - use fal_client::requests::webhook::image::edit::enqueue_gemini_25_flash_edit_webhook::Gemini25FlashEditAspectRatio; - use fal_client::requests::webhook::image::text::enqueue_gemini_25_flash_text_to_image_webhook::Gemini25FlashTextToImageAspectRatio; - - /// Distill a text request and run an assertion against the inner plan. - fn with_text_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill_text(request); - match distilled.plan() { - ImageGenerationPlan::FalNanoBanana(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalNanoBanana, got {:?}", other), - } - } - - /// Distill an edit request and run an assertion against the inner plan. - fn with_edit_plan( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - assertion: F, - ) { - let distilled = distill_edit(request, hydration); - match distilled.plan() { - ImageGenerationPlan::FalNanoBanana(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalNanoBanana, got {:?}", other), - } - } - - mod text { - use super::*; - - fn assert_t2i_for_aspect_ratio( - ar: Option, - assertion: F, - ) { - with_text_plan(&make_request(Some("p"), ar, Some(1), None), assertion); - } - - fn assert_t2i_for_batch( - batch: Option, - assertion: F, - ) { - with_text_plan(&make_request(Some("p"), None, batch, None), assertion); - } - - // ── Mode detection ──────────────────────────────────────────────────── - - #[test] - fn text_mode_has_empty_image_urls() { - with_text_plan(&make_request(Some("p"), None, Some(1), None), |plan| { - assert!(plan.image_urls.is_empty(), "text mode must have no image_urls"); - }); - } - - // ── Aspect ratio direct mappings ────────────────────────────────────── - - #[test] - fn default_aspect_ratio_is_none() { - // No aspect ratio specified → t2i_aspect_ratio is None. - assert_t2i_for_aspect_ratio(None, |plan| { - assert!(plan.t2i_aspect_ratio.is_none()); - }); - } - - #[test] - fn square_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn square_hd_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn wide_5x4_yields_five_by_four() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::FiveByFour) - )); - }); - } - - #[test] - fn wide_4x3_yields_four_by_three() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::FourByThree) - )); - }); - } - - #[test] - fn wide_3x2_yields_three_by_two() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideThreeByTwo), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::ThreeByTwo) - )); - }); - } - - #[test] - fn wide_16x9_yields_sixteen_by_nine() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_alias_yields_sixteen_by_nine() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_21x9_yields_twenty_one_by_nine() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideTwentyOneByNine), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::TwentyOneByNine) - )); - }); - } - - #[test] - fn tall_4x5_yields_four_by_five() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallFourByFive), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::FourByFive) - )); - }); - } - - #[test] - fn tall_3x4_yields_three_by_four() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::ThreeByFour) - )); - }); - } - - #[test] - fn tall_2x3_yields_two_by_three() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::TwoByThree) - )); - }); - } - - #[test] - fn tall_9x16_yields_nine_by_sixteen() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::NineBySixteen) - )); - }); - } - - #[test] - fn tall_alias_yields_nine_by_sixteen() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::NineBySixteen) - )); - }); - } - - // ── Auto* fallbacks (Auto is not valid for t2i; falls back to square) - - #[test] - fn auto_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn auto_2k_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Auto2k), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn auto_4k_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Auto4k), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(Gemini25FlashTextToImageAspectRatio::OneByOne) - )); - }); - } - - // ── Num images mapping ──────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - assert_t2i_for_batch(None, |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::One)); - }); - } - - #[test] - fn batch_of_one_yields_one() { - assert_t2i_for_batch(Some(1), |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::One)); - }); - } - - #[test] - fn batch_of_two_yields_two() { - assert_t2i_for_batch(Some(2), |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::Two)); - }); - } - - #[test] - fn batch_of_three_yields_three() { - assert_t2i_for_batch(Some(3), |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::Three)); - }); - } - - #[test] - fn batch_of_four_yields_four() { - assert_t2i_for_batch(Some(4), |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::Four)); - }); - } - - #[test] - fn batch_above_four_clamps_to_four() { - assert_t2i_for_batch(Some(9), |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_text_plan(&make_request(Some("a corgi in a hat"), None, Some(1), None), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_text_plan(&make_request(None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } - - mod edit { - use super::*; - - fn assert_edit_for_aspect_ratio( - ar: Option, - assertion: F, - ) { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), ar, Some(1), Some(tokens)), - &hydration, - assertion, - ); - } - - fn assert_edit_for_batch( - batch: Option, - assertion: F, - ) { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, batch, Some(tokens)), - &hydration, - assertion, - ); - } - - // ── Mode detection / image-url passthrough ──────────────────────────── - - #[test] - fn edit_mode_populates_image_urls() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 1); - assert!(plan.image_urls[0].starts_with("https://fake.example.com/")); - }, - ); - } - - #[test] - fn edit_mode_with_two_image_refs() { - let (tokens, hydration) = fake_image_refs(2); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 2); - }, - ); - } - - #[test] - fn edit_mode_with_five_image_refs() { - let (tokens, hydration) = fake_image_refs(5); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 5); - }, - ); - } - - #[test] - fn edit_image_urls_are_hydrated_from_map() { - let (tokens, hydration) = fake_image_refs(3); - let expected: Vec = tokens - .iter() - .map(|t| hydration.get(t).unwrap().to_string()) - .collect(); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens.clone())), - &hydration, - |plan| { - // Order should be preserved (token order in → URL order out). - assert_eq!(plan.image_urls, expected); - }, - ); - } - - // ── Aspect ratio direct mappings ────────────────────────────────────── - - #[test] - fn default_aspect_ratio_is_none() { - assert_edit_for_aspect_ratio(None, |plan| { - assert!(plan.edit_aspect_ratio.is_none()); - }); - } - - #[test] - fn square_yields_one_by_one() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn square_hd_yields_one_by_one() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn wide_5x4_yields_five_by_four() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::FiveByFour) - )); - }); - } - - #[test] - fn wide_4x3_yields_four_by_three() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::FourByThree) - )); - }); - } - - #[test] - fn wide_3x2_yields_three_by_two() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideThreeByTwo), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::ThreeByTwo) - )); - }); - } - - #[test] - fn wide_16x9_yields_sixteen_by_nine() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_alias_yields_sixteen_by_nine() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_21x9_yields_twenty_one_by_nine() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideTwentyOneByNine), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::TwentyOneByNine) - )); - }); - } - - #[test] - fn tall_4x5_yields_four_by_five() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallFourByFive), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::FourByFive) - )); - }); - } - - #[test] - fn tall_3x4_yields_three_by_four() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::ThreeByFour) - )); - }); - } - - #[test] - fn tall_2x3_yields_two_by_three() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::TwoByThree) - )); - }); - } - - #[test] - fn tall_9x16_yields_nine_by_sixteen() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::NineBySixteen) - )); - }); - } - - #[test] - fn tall_alias_yields_nine_by_sixteen() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::NineBySixteen) - )); - }); - } - - // ── Auto* maps to Auto in edit mode (unlike t2i) ────────────────────── - - #[test] - fn auto_yields_auto() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::Auto) - )); - }); - } - - #[test] - fn auto_2k_yields_auto() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Auto2k), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::Auto) - )); - }); - } - - #[test] - fn auto_4k_yields_auto() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Auto4k), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(Gemini25FlashEditAspectRatio::Auto) - )); - }); - } - - // ── Num images mapping ──────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - assert_edit_for_batch(None, |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::One)); - }); - } - - #[test] - fn batch_of_one_yields_one() { - assert_edit_for_batch(Some(1), |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::One)); - }); - } - - #[test] - fn batch_of_two_yields_two() { - assert_edit_for_batch(Some(2), |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::Two)); - }); - } - - #[test] - fn batch_of_three_yields_three() { - assert_edit_for_batch(Some(3), |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::Three)); - }); - } - - #[test] - fn batch_of_four_yields_four() { - assert_edit_for_batch(Some(4), |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::Four)); - }); - } - - #[test] - fn batch_above_four_clamps_to_four() { - assert_edit_for_batch(Some(9), |plan| { - assert!(matches!(plan.num_images, FalNbNumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("make it shiny"), None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.prompt, Some("make it shiny".to_string())); - }, - ); - } - - #[test] - fn missing_prompt_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.prompt, None); - }, - ); - } - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_2.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_2.rs deleted file mode 100644 index 473e0839b4..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_2.rs +++ /dev/null @@ -1,1321 +0,0 @@ -//! Tests for omni-gen against the Nano Banana 2 model. -//! -//! Like Nano Banana, this model is multi-function: text-to-image and -//! image-to-image (edit). It additionally supports a `resolution` parameter -//! (0.5K / 1K / 2K / 4K) that affects pricing — 4K outputs cost roughly 2× -//! the 1K/2K rate, and 0.5K is the discount tier. -//! -//! Artcraft-tier pricing (the rate the user is billed) is scaled off the 1K -//! base cost of $0.08: -//! 0.5K → $0.06 / image (0.75×) -//! 1K → $0.08 / image (1×, default) -//! 2K → $0.12 / image (1.5×) -//! 4K → $0.16 / image (2×) -//! Same rates apply to both text-to-image and image-edit modes. - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use enums::common::generation::common_resolution::CommonResolution; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - /// Build a minimal Nano Banana 2 omni request. - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - resolution: Option, - image_batch_count: Option, - image_media_tokens: Option>, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::NanoBanana2), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens, - resolution, - aspect_ratio, - quality: None, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - /// Make `n` fake media tokens paired with fake URLs. - fn fake_image_refs(n: usize) -> (Vec, HashMap) { - let mut tokens = Vec::with_capacity(n); - let mut map = HashMap::with_capacity(n); - for i in 0..n { - let token = MediaFileToken::new_from_str(&format!("mf_test{:028}", i)); - let url = Url::parse(&format!("https://fake.example.com/img{}.png", i)).unwrap(); - map.insert(token.clone(), url); - tokens.push(token); - } - (tokens, map) - } - - /// Distill a text-to-image request (no image refs, empty hydration map). - fn distill_text(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for Nano Banana 2 (text)") - } - - /// Distill an edit request, providing the hydration map. - fn distill_edit( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - ) -> DistilledImageRequest { - distill_image_request(request, Some(hydration)) - .expect("distill_image_request should succeed for Nano Banana 2 (edit)") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST - // ──────────────────────────────────────────────────────────────────────────── - // - // Pricing comes from `estimate_image_cost_artcraft_nano_banana_2`: - // 0.5K → 6¢/image, 1K → 8¢/image (default), 2K → 12¢/image, 4K → 16¢/image. - // - // Same rates apply to both text-to-image and edit modes. - mod cost { - use super::*; - - mod text { - use super::*; - - fn cost_for( - resolution: Option, - image_batch_count: Option, - ) -> (Option, Option) { - let request = make_request(Some("a happy puppy"), None, resolution, image_batch_count, None); - let distilled = distill_text(&request); - ( - distilled.cost.cost_in_credits, - distilled.cost.cost_in_usd_cents, - ) - } - - // ── Default resolution (None → 1K → 8¢/image) ────────────────────────── - - #[test] - fn default_resolution_default_batch_costs_8_cents() { - let (credits, cents) = cost_for(None, None); - assert_eq!(credits, Some(8)); - assert_eq!(cents, Some(8)); - } - - #[test] - fn default_resolution_batch_of_one_costs_8_cents() { - let (_c, cents) = cost_for(None, Some(1)); - assert_eq!(cents, Some(8)); - } - - #[test] - fn default_resolution_batch_of_two_costs_16_cents() { - let (_c, cents) = cost_for(None, Some(2)); - assert_eq!(cents, Some(16)); - } - - #[test] - fn default_resolution_batch_of_three_costs_24_cents() { - let (_c, cents) = cost_for(None, Some(3)); - assert_eq!(cents, Some(24)); - } - - #[test] - fn default_resolution_batch_of_four_costs_32_cents() { - let (_c, cents) = cost_for(None, Some(4)); - assert_eq!(cents, Some(32)); - } - - #[test] - fn default_resolution_batch_above_max_clamps_to_32_cents() { - let (_c, cents) = cost_for(None, Some(7)); - assert_eq!(cents, Some(32)); - } - - // ── 0.5K resolution (6¢/image, 0.75× base) ───────────────────────────── - - #[test] - fn half_k_batch_of_one_costs_6_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(1)); - assert_eq!(cents, Some(6)); - } - - #[test] - fn half_k_batch_of_two_costs_12_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(2)); - assert_eq!(cents, Some(12)); - } - - #[test] - fn half_k_batch_of_three_costs_18_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(3)); - assert_eq!(cents, Some(18)); - } - - #[test] - fn half_k_batch_of_four_costs_24_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(4)); - assert_eq!(cents, Some(24)); - } - - #[test] - fn half_k_batch_above_max_clamps_to_24_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(9)); - assert_eq!(cents, Some(24)); - } - - // ── 1K resolution (8¢/image, base) ───────────────────────────────────── - - #[test] - fn one_k_batch_of_one_costs_8_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(1)); - assert_eq!(cents, Some(8)); - } - - #[test] - fn one_k_batch_of_two_costs_16_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(2)); - assert_eq!(cents, Some(16)); - } - - #[test] - fn one_k_batch_of_three_costs_24_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(3)); - assert_eq!(cents, Some(24)); - } - - #[test] - fn one_k_batch_of_four_costs_32_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(4)); - assert_eq!(cents, Some(32)); - } - - #[test] - fn one_k_batch_above_max_clamps_to_32_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(9)); - assert_eq!(cents, Some(32)); - } - - // ── 2K resolution (12¢/image, 1.5× base) ─────────────────────────────── - - #[test] - fn two_k_batch_of_one_costs_12_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(1)); - assert_eq!(cents, Some(12)); - } - - #[test] - fn two_k_batch_of_two_costs_24_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(2)); - assert_eq!(cents, Some(24)); - } - - #[test] - fn two_k_batch_of_three_costs_36_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(3)); - assert_eq!(cents, Some(36)); - } - - #[test] - fn two_k_batch_of_four_costs_48_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(4)); - assert_eq!(cents, Some(48)); - } - - #[test] - fn two_k_batch_above_max_clamps_to_48_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(9)); - assert_eq!(cents, Some(48)); - } - - // ── 4K resolution (16¢/image, 2× base) ───────────────────────────────── - - #[test] - fn four_k_batch_of_one_costs_16_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(1)); - assert_eq!(cents, Some(16)); - } - - #[test] - fn four_k_batch_of_two_costs_32_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(2)); - assert_eq!(cents, Some(32)); - } - - #[test] - fn four_k_batch_of_three_costs_48_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(3)); - assert_eq!(cents, Some(48)); - } - - #[test] - fn four_k_batch_of_four_costs_64_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(4)); - assert_eq!(cents, Some(64)); - } - - #[test] - fn four_k_batch_above_max_clamps_to_64_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(9)); - assert_eq!(cents, Some(64)); - } - - // ── 3K (unsupported) falls back to 2K pricing (= 12¢/image) ──────────── - - #[test] - fn three_k_one_image_falls_back_to_two_k_pricing() { - let (_c, cents) = cost_for(Some(CommonResolution::ThreeK), Some(1)); - assert_eq!(cents, Some(12)); - } - - #[test] - fn three_k_four_images_falls_back_to_two_k_pricing() { - let (_c, cents) = cost_for(Some(CommonResolution::ThreeK), Some(4)); - assert_eq!(cents, Some(48)); - } - - // ── Misc independence checks ─────────────────────────────────────────── - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let aspect_ratios = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::SquareHd), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::TallNineBySixteen), - Some(CommonAspectRatio::Auto), - ]; - for ar in aspect_ratios { - let request = make_request(Some("p"), ar, None, Some(2), None); - let distilled = distill_text(&request); - assert_eq!( - distilled.cost.cost_in_usd_cents, - Some(16), - "expected 16¢ regardless of aspect ratio (got {:?})", - ar, - ); - } - } - - #[test] - fn cost_is_independent_of_prompt() { - let with_prompt = distill_text(&make_request(Some("a cat"), None, None, Some(3), None)); - let without_prompt = distill_text(&make_request(None, None, None, Some(3), None)); - assert_eq!(with_prompt.cost.cost_in_usd_cents, Some(24)); - assert_eq!(without_prompt.cost.cost_in_usd_cents, Some(24)); - } - - #[test] - fn cost_metadata_flags_are_default() { - let request = make_request(Some("p"), None, None, Some(1), None); - let distilled = distill_text(&request); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - } - - mod edit { - use super::*; - - fn cost_for( - resolution: Option, - image_batch_count: Option, - num_image_refs: usize, - ) -> (Option, Option) { - let (tokens, hydration) = fake_image_refs(num_image_refs); - let request = make_request( - Some("make it pop"), - None, - resolution, - image_batch_count, - Some(tokens), - ); - let distilled = distill_edit(&request, &hydration); - ( - distilled.cost.cost_in_credits, - distilled.cost.cost_in_usd_cents, - ) - } - - // ── Default resolution (None → 1K → 8¢/image) ────────────────────────── - - #[test] - fn default_resolution_default_batch_costs_8_cents() { - let (credits, cents) = cost_for(None, None, 1); - assert_eq!(credits, Some(8)); - assert_eq!(cents, Some(8)); - } - - #[test] - fn default_resolution_batch_of_one_costs_8_cents() { - let (_c, cents) = cost_for(None, Some(1), 1); - assert_eq!(cents, Some(8)); - } - - #[test] - fn default_resolution_batch_of_two_costs_16_cents() { - let (_c, cents) = cost_for(None, Some(2), 1); - assert_eq!(cents, Some(16)); - } - - #[test] - fn default_resolution_batch_of_three_costs_24_cents() { - let (_c, cents) = cost_for(None, Some(3), 1); - assert_eq!(cents, Some(24)); - } - - #[test] - fn default_resolution_batch_of_four_costs_32_cents() { - let (_c, cents) = cost_for(None, Some(4), 1); - assert_eq!(cents, Some(32)); - } - - #[test] - fn default_resolution_batch_above_max_clamps_to_32_cents() { - let (_c, cents) = cost_for(None, Some(9), 1); - assert_eq!(cents, Some(32)); - } - - // ── 0.5K (6¢/image) ──────────────────────────────────────────────────── - - #[test] - fn half_k_batch_of_one_costs_6_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(1), 1); - assert_eq!(cents, Some(6)); - } - - #[test] - fn half_k_batch_of_two_costs_12_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(2), 1); - assert_eq!(cents, Some(12)); - } - - #[test] - fn half_k_batch_of_three_costs_18_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(3), 1); - assert_eq!(cents, Some(18)); - } - - #[test] - fn half_k_batch_of_four_costs_24_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(4), 1); - assert_eq!(cents, Some(24)); - } - - // ── 1K (8¢/image, base) ──────────────────────────────────────────────── - - #[test] - fn one_k_batch_of_one_costs_8_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(1), 1); - assert_eq!(cents, Some(8)); - } - - #[test] - fn one_k_batch_of_two_costs_16_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(2), 1); - assert_eq!(cents, Some(16)); - } - - #[test] - fn one_k_batch_of_four_costs_32_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(4), 1); - assert_eq!(cents, Some(32)); - } - - // ── 2K (12¢/image, 1.5× base) ────────────────────────────────────────── - - #[test] - fn two_k_batch_of_one_costs_12_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(1), 1); - assert_eq!(cents, Some(12)); - } - - #[test] - fn two_k_batch_of_two_costs_24_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(2), 1); - assert_eq!(cents, Some(24)); - } - - #[test] - fn two_k_batch_of_three_costs_36_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(3), 1); - assert_eq!(cents, Some(36)); - } - - #[test] - fn two_k_batch_of_four_costs_48_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(4), 1); - assert_eq!(cents, Some(48)); - } - - // ── 4K (16¢/image, 2× base) ──────────────────────────────────────────── - - #[test] - fn four_k_batch_of_one_costs_16_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(1), 1); - assert_eq!(cents, Some(16)); - } - - #[test] - fn four_k_batch_of_two_costs_32_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(2), 1); - assert_eq!(cents, Some(32)); - } - - #[test] - fn four_k_batch_of_three_costs_48_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(3), 1); - assert_eq!(cents, Some(48)); - } - - #[test] - fn four_k_batch_of_four_costs_64_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(4), 1); - assert_eq!(cents, Some(64)); - } - - // ── 3K (unsupported) falls back to 2K pricing ────────────────────────── - - #[test] - fn three_k_one_image_falls_back_to_two_k_pricing() { - let (_c, cents) = cost_for(Some(CommonResolution::ThreeK), Some(1), 1); - assert_eq!(cents, Some(12)); - } - - #[test] - fn three_k_four_images_falls_back_to_two_k_pricing() { - let (_c, cents) = cost_for(Some(CommonResolution::ThreeK), Some(4), 1); - assert_eq!(cents, Some(48)); - } - - #[test] - fn cost_is_independent_of_image_ref_count() { - for num_refs in [1usize, 2, 3, 5] { - let (_c, cents) = cost_for(None, Some(2), num_refs); - assert_eq!( - cents, - Some(16), - "expected 16¢ regardless of {} image refs", - num_refs, - ); - } - } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let aspect_ratios = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::Auto), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::TallNineBySixteen), - ]; - for ar in aspect_ratios { - let (tokens, hydration) = fake_image_refs(2); - let request = make_request(Some("p"), ar, None, Some(2), Some(tokens)); - let distilled = distill_edit(&request, &hydration); - assert_eq!( - distilled.cost.cost_in_usd_cents, - Some(16), - "expected 16¢ regardless of aspect ratio (got {:?})", - ar, - ); - } - } - - #[test] - fn cost_metadata_flags_are_default() { - let (tokens, hydration) = fake_image_refs(1); - let request = make_request(Some("p"), None, None, Some(1), Some(tokens)); - let distilled = distill_edit(&request, &hydration); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - - #[test] - fn edit_and_text_cost_match_at_each_resolution() { - // Both modes are billed at the same per-resolution rate. - let resolutions = [ - None, - Some(CommonResolution::HalfK), - Some(CommonResolution::OneK), - Some(CommonResolution::TwoK), - Some(CommonResolution::FourK), - ]; - for res in resolutions { - for batch in [1u16, 2, 3, 4] { - let text = distill_text(&make_request(Some("p"), None, res, Some(batch), None)); - let (tokens, hydration) = fake_image_refs(1); - let edit = distill_edit( - &make_request(Some("p"), None, res, Some(batch), Some(tokens)), - &hydration, - ); - assert_eq!( - text.cost.cost_in_usd_cents, - edit.cost.cost_in_usd_cents, - "text/edit cost diverged at res={:?} batch={}", - res, - batch, - ); - } - } - } - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - // - // The distilled plan is `ImageGenerationPlan::FalNanaBanana2(PlanFalNanaBanana2)`. - // It carries: - // - prompt: passthrough - // - image_urls: empty in text mode, populated in edit mode - // - t2i_aspect_ratio / edit_aspect_ratio: pre-resolved per mode - // - resolution: shared between modes - // - num_images: 1..=4 (clamped under PayMoreUpgrade) - // - // The handler picks t2i vs edit based on whether image_urls is empty. - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana_2::{ - FalNb2NumImages, FalNb2Resolution, PlanFalNanaBanana2, - }; - use fal_client::requests::webhook::image::edit::enqueue_nano_banana_2_edit_image_webhook::EnqueueNanoBanana2EditImageAspectRatio; - use fal_client::requests::webhook::image::text::enqueue_nano_banana_2_text_to_image_webhook::EnqueueNanoBanana2TextToImageAspectRatio; - - fn with_text_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill_text(request); - match distilled.plan() { - ImageGenerationPlan::FalNanaBanana2(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalNanaBanana2, got {:?}", other), - } - } - - fn with_edit_plan( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - assertion: F, - ) { - let distilled = distill_edit(request, hydration); - match distilled.plan() { - ImageGenerationPlan::FalNanaBanana2(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalNanaBanana2, got {:?}", other), - } - } - - mod text { - use super::*; - - fn assert_t2i_for_aspect_ratio( - ar: Option, - assertion: F, - ) { - with_text_plan(&make_request(Some("p"), ar, None, Some(1), None), assertion); - } - - fn assert_t2i_for_batch( - batch: Option, - assertion: F, - ) { - with_text_plan(&make_request(Some("p"), None, None, batch, None), assertion); - } - - fn assert_t2i_for_resolution( - res: Option, - assertion: F, - ) { - with_text_plan(&make_request(Some("p"), None, res, Some(1), None), assertion); - } - - // ── Mode detection ──────────────────────────────────────────────────── - - #[test] - fn text_mode_has_empty_image_urls() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(plan.image_urls.is_empty(), "text mode must have no image_urls"); - }); - } - - // ── Aspect ratio direct mappings ────────────────────────────────────── - - #[test] - fn default_aspect_ratio_is_none() { - assert_t2i_for_aspect_ratio(None, |plan| { - assert!(plan.t2i_aspect_ratio.is_none()); - }); - } - - #[test] - fn square_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn square_hd_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn wide_5x4_yields_five_by_four() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::FiveByFour) - )); - }); - } - - #[test] - fn wide_4x3_yields_four_by_three() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::FourByThree) - )); - }); - } - - #[test] - fn wide_3x2_yields_three_by_two() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideThreeByTwo), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::ThreeByTwo) - )); - }); - } - - #[test] - fn wide_16x9_yields_sixteen_by_nine() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_alias_yields_sixteen_by_nine() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_21x9_yields_twenty_one_by_nine() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideTwentyOneByNine), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::TwentyOneByNine) - )); - }); - } - - #[test] - fn tall_4x5_yields_four_by_five() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallFourByFive), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::FourByFive) - )); - }); - } - - #[test] - fn tall_3x4_yields_three_by_four() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::ThreeByFour) - )); - }); - } - - #[test] - fn tall_2x3_yields_two_by_three() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::TwoByThree) - )); - }); - } - - #[test] - fn tall_9x16_yields_nine_by_sixteen() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::NineBySixteen) - )); - }); - } - - #[test] - fn tall_alias_yields_nine_by_sixteen() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::NineBySixteen) - )); - }); - } - - // ── Auto* maps to Auto in t2i for NB2 (unlike NB1) ──────────────────── - - #[test] - fn auto_yields_auto() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::Auto) - )); - }); - } - - #[test] - fn auto_2k_yields_auto() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Auto2k), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::Auto) - )); - }); - } - - #[test] - fn auto_4k_yields_auto() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Auto4k), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBanana2TextToImageAspectRatio::Auto) - )); - }); - } - - // ── Resolution mapping ──────────────────────────────────────────────── - - #[test] - fn default_resolution_is_none() { - assert_t2i_for_resolution(None, |plan| { - assert!(plan.resolution.is_none()); - }); - } - - #[test] - fn half_k_yields_half_k() { - assert_t2i_for_resolution(Some(CommonResolution::HalfK), |plan| { - assert!(matches!(plan.resolution, Some(FalNb2Resolution::HalfK))); - }); - } - - #[test] - fn one_k_yields_one_k() { - assert_t2i_for_resolution(Some(CommonResolution::OneK), |plan| { - assert!(matches!(plan.resolution, Some(FalNb2Resolution::OneK))); - }); - } - - #[test] - fn two_k_yields_two_k() { - assert_t2i_for_resolution(Some(CommonResolution::TwoK), |plan| { - assert!(matches!(plan.resolution, Some(FalNb2Resolution::TwoK))); - }); - } - - #[test] - fn four_k_yields_four_k() { - assert_t2i_for_resolution(Some(CommonResolution::FourK), |plan| { - assert!(matches!(plan.resolution, Some(FalNb2Resolution::FourK))); - }); - } - - #[test] - fn three_k_falls_back_to_two_k() { - assert_t2i_for_resolution(Some(CommonResolution::ThreeK), |plan| { - assert!(matches!(plan.resolution, Some(FalNb2Resolution::TwoK))); - }); - } - - // ── Num images mapping ──────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - assert_t2i_for_batch(None, |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::One)); - }); - } - - #[test] - fn batch_of_one_yields_one() { - assert_t2i_for_batch(Some(1), |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::One)); - }); - } - - #[test] - fn batch_of_two_yields_two() { - assert_t2i_for_batch(Some(2), |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::Two)); - }); - } - - #[test] - fn batch_of_three_yields_three() { - assert_t2i_for_batch(Some(3), |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::Three)); - }); - } - - #[test] - fn batch_of_four_yields_four() { - assert_t2i_for_batch(Some(4), |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::Four)); - }); - } - - #[test] - fn batch_above_four_clamps_to_four() { - assert_t2i_for_batch(Some(9), |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_text_plan(&make_request(Some("a corgi in a hat"), None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_text_plan(&make_request(None, None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } - - mod edit { - use super::*; - - fn assert_edit_for_aspect_ratio( - ar: Option, - assertion: F, - ) { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), ar, None, Some(1), Some(tokens)), - &hydration, - assertion, - ); - } - - fn assert_edit_for_batch( - batch: Option, - assertion: F, - ) { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, batch, Some(tokens)), - &hydration, - assertion, - ); - } - - fn assert_edit_for_resolution( - res: Option, - assertion: F, - ) { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, res, Some(1), Some(tokens)), - &hydration, - assertion, - ); - } - - // ── Mode detection / image-url passthrough ──────────────────────────── - - #[test] - fn edit_mode_populates_image_urls() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 1); - assert!(plan.image_urls[0].starts_with("https://fake.example.com/")); - }, - ); - } - - #[test] - fn edit_mode_with_two_image_refs() { - let (tokens, hydration) = fake_image_refs(2); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 2); - }, - ); - } - - #[test] - fn edit_mode_with_five_image_refs() { - let (tokens, hydration) = fake_image_refs(5); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 5); - }, - ); - } - - #[test] - fn edit_image_urls_are_hydrated_from_map() { - let (tokens, hydration) = fake_image_refs(3); - let expected: Vec = tokens - .iter() - .map(|t| hydration.get(t).unwrap().to_string()) - .collect(); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens.clone())), - &hydration, - |plan| { - assert_eq!(plan.image_urls, expected); - }, - ); - } - - // ── Aspect ratio direct mappings ────────────────────────────────────── - - #[test] - fn default_aspect_ratio_is_none() { - assert_edit_for_aspect_ratio(None, |plan| { - assert!(plan.edit_aspect_ratio.is_none()); - }); - } - - #[test] - fn square_yields_one_by_one() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn square_hd_yields_one_by_one() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn wide_5x4_yields_five_by_four() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::FiveByFour) - )); - }); - } - - #[test] - fn wide_4x3_yields_four_by_three() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::FourByThree) - )); - }); - } - - #[test] - fn wide_3x2_yields_three_by_two() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideThreeByTwo), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::ThreeByTwo) - )); - }); - } - - #[test] - fn wide_16x9_yields_sixteen_by_nine() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_alias_yields_sixteen_by_nine() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_21x9_yields_twenty_one_by_nine() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideTwentyOneByNine), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::TwentyOneByNine) - )); - }); - } - - #[test] - fn tall_4x5_yields_four_by_five() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallFourByFive), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::FourByFive) - )); - }); - } - - #[test] - fn tall_3x4_yields_three_by_four() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::ThreeByFour) - )); - }); - } - - #[test] - fn tall_2x3_yields_two_by_three() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::TwoByThree) - )); - }); - } - - #[test] - fn tall_9x16_yields_nine_by_sixteen() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::NineBySixteen) - )); - }); - } - - #[test] - fn tall_alias_yields_nine_by_sixteen() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::NineBySixteen) - )); - }); - } - - // ── Auto* maps to Auto in edit mode ─────────────────────────────────── - - #[test] - fn auto_yields_auto() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::Auto) - )); - }); - } - - #[test] - fn auto_2k_yields_auto() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Auto2k), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::Auto) - )); - }); - } - - #[test] - fn auto_4k_yields_auto() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Auto4k), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBanana2EditImageAspectRatio::Auto) - )); - }); - } - - // ── Resolution mapping ──────────────────────────────────────────────── - - #[test] - fn default_resolution_is_none() { - assert_edit_for_resolution(None, |plan| { - assert!(plan.resolution.is_none()); - }); - } - - #[test] - fn half_k_yields_half_k() { - assert_edit_for_resolution(Some(CommonResolution::HalfK), |plan| { - assert!(matches!(plan.resolution, Some(FalNb2Resolution::HalfK))); - }); - } - - #[test] - fn one_k_yields_one_k() { - assert_edit_for_resolution(Some(CommonResolution::OneK), |plan| { - assert!(matches!(plan.resolution, Some(FalNb2Resolution::OneK))); - }); - } - - #[test] - fn two_k_yields_two_k() { - assert_edit_for_resolution(Some(CommonResolution::TwoK), |plan| { - assert!(matches!(plan.resolution, Some(FalNb2Resolution::TwoK))); - }); - } - - #[test] - fn four_k_yields_four_k() { - assert_edit_for_resolution(Some(CommonResolution::FourK), |plan| { - assert!(matches!(plan.resolution, Some(FalNb2Resolution::FourK))); - }); - } - - #[test] - fn three_k_falls_back_to_two_k() { - assert_edit_for_resolution(Some(CommonResolution::ThreeK), |plan| { - assert!(matches!(plan.resolution, Some(FalNb2Resolution::TwoK))); - }); - } - - // ── Num images mapping ──────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - assert_edit_for_batch(None, |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::One)); - }); - } - - #[test] - fn batch_of_one_yields_one() { - assert_edit_for_batch(Some(1), |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::One)); - }); - } - - #[test] - fn batch_of_two_yields_two() { - assert_edit_for_batch(Some(2), |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::Two)); - }); - } - - #[test] - fn batch_of_three_yields_three() { - assert_edit_for_batch(Some(3), |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::Three)); - }); - } - - #[test] - fn batch_of_four_yields_four() { - assert_edit_for_batch(Some(4), |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::Four)); - }); - } - - #[test] - fn batch_above_four_clamps_to_four() { - assert_edit_for_batch(Some(9), |plan| { - assert!(matches!(plan.num_images, FalNb2NumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("make it shiny"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.prompt, Some("make it shiny".to_string())); - }, - ); - } - - #[test] - fn missing_prompt_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(None, None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.prompt, None); - }, - ); - } - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_pro.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_pro.rs deleted file mode 100644 index 9fa822c33e..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_pro.rs +++ /dev/null @@ -1,1211 +0,0 @@ -//! Tests for omni-gen against the Nano Banana Pro model. -//! -//! Like Nano Banana / Nano Banana 2, this model is multi-function (text-to-image -//! and image-to-image edit) and supports a `resolution` parameter (1K / 2K / 4K) -//! that affects pricing. Unlike NB2, Pro has no native 0.5K tier — sub-1K -//! `CommonResolution` values fall back to 1K. -//! -//! Artcraft-tier pricing (the rate the user is billed), matching the legacy -//! `nano_banana_pro_multi_function_image_gen_handler` and the upstream Fal cost: -//! 1K → $0.15 / image (default) -//! 2K → $0.15 / image -//! 4K → $0.30 / image (2× the standard rate) -//! Same rates apply to both text-to-image and image-edit modes. - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use enums::common::generation::common_resolution::CommonResolution; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - resolution: Option, - image_batch_count: Option, - image_media_tokens: Option>, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::NanoBananaPro), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens, - resolution, - aspect_ratio, - quality: None, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - fn fake_image_refs(n: usize) -> (Vec, HashMap) { - let mut tokens = Vec::with_capacity(n); - let mut map = HashMap::with_capacity(n); - for i in 0..n { - let token = MediaFileToken::new_from_str(&format!("mf_test{:028}", i)); - let url = Url::parse(&format!("https://fake.example.com/img{}.png", i)).unwrap(); - map.insert(token.clone(), url); - tokens.push(token); - } - (tokens, map) - } - - fn distill_text(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for Nano Banana Pro (text)") - } - - fn distill_edit( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - ) -> DistilledImageRequest { - distill_image_request(request, Some(hydration)) - .expect("distill_image_request should succeed for Nano Banana Pro (edit)") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST - // ──────────────────────────────────────────────────────────────────────────── - // - // Pricing comes from `estimate_image_cost_artcraft_nano_banana_pro`: - // 1K → 15¢/image (default), 2K → 15¢/image, 4K → 30¢/image. - // - // Same rates apply to both text-to-image and edit modes. - mod cost { - use super::*; - - mod text { - use super::*; - - fn cost_for( - resolution: Option, - image_batch_count: Option, - ) -> (Option, Option) { - let request = make_request(Some("a happy puppy"), None, resolution, image_batch_count, None); - let distilled = distill_text(&request); - ( - distilled.cost.cost_in_credits, - distilled.cost.cost_in_usd_cents, - ) - } - - // ── Default resolution (None → 1K → 15¢/image) ───────────────────────── - - #[test] - fn default_resolution_default_batch_costs_15_cents() { - let (credits, cents) = cost_for(None, None); - assert_eq!(credits, Some(15)); - assert_eq!(cents, Some(15)); - } - - #[test] - fn default_resolution_batch_of_one_costs_15_cents() { - let (_c, cents) = cost_for(None, Some(1)); - assert_eq!(cents, Some(15)); - } - - #[test] - fn default_resolution_batch_of_two_costs_30_cents() { - let (_c, cents) = cost_for(None, Some(2)); - assert_eq!(cents, Some(30)); - } - - #[test] - fn default_resolution_batch_of_three_costs_45_cents() { - let (_c, cents) = cost_for(None, Some(3)); - assert_eq!(cents, Some(45)); - } - - #[test] - fn default_resolution_batch_of_four_costs_60_cents() { - let (_c, cents) = cost_for(None, Some(4)); - assert_eq!(cents, Some(60)); - } - - #[test] - fn default_resolution_batch_above_max_clamps_to_60_cents() { - let (_c, cents) = cost_for(None, Some(7)); - assert_eq!(cents, Some(60)); - } - - // ── 1K resolution (15¢/image) ────────────────────────────────────────── - - #[test] - fn one_k_batch_of_one_costs_15_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(1)); - assert_eq!(cents, Some(15)); - } - - #[test] - fn one_k_batch_of_two_costs_30_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(2)); - assert_eq!(cents, Some(30)); - } - - #[test] - fn one_k_batch_of_three_costs_45_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(3)); - assert_eq!(cents, Some(45)); - } - - #[test] - fn one_k_batch_of_four_costs_60_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(4)); - assert_eq!(cents, Some(60)); - } - - // ── 2K resolution (15¢/image, same as 1K) ────────────────────────────── - - #[test] - fn two_k_batch_of_one_costs_15_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(1)); - assert_eq!(cents, Some(15)); - } - - #[test] - fn two_k_batch_of_two_costs_30_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(2)); - assert_eq!(cents, Some(30)); - } - - #[test] - fn two_k_batch_of_three_costs_45_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(3)); - assert_eq!(cents, Some(45)); - } - - #[test] - fn two_k_batch_of_four_costs_60_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(4)); - assert_eq!(cents, Some(60)); - } - - // ── 4K resolution (30¢/image, 2× standard rate) ──────────────────────── - - #[test] - fn four_k_batch_of_one_costs_30_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(1)); - assert_eq!(cents, Some(30)); - } - - #[test] - fn four_k_batch_of_two_costs_60_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(2)); - assert_eq!(cents, Some(60)); - } - - #[test] - fn four_k_batch_of_three_costs_90_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(3)); - assert_eq!(cents, Some(90)); - } - - #[test] - fn four_k_batch_of_four_costs_120_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(4)); - assert_eq!(cents, Some(120)); - } - - // ── 0.5K and 3K (unsupported) fall back to other tiers ───────────────── - - #[test] - fn half_k_falls_back_to_one_k_pricing() { - // Pro has no 0.5K tier; nearest supported is 1K → 15¢/image. - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(1)); - assert_eq!(cents, Some(15)); - } - - #[test] - fn three_k_falls_back_to_two_k_pricing() { - let (_c, cents) = cost_for(Some(CommonResolution::ThreeK), Some(1)); - assert_eq!(cents, Some(15)); - } - - // ── Misc independence checks ─────────────────────────────────────────── - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let aspect_ratios = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::SquareHd), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::TallNineBySixteen), - Some(CommonAspectRatio::Auto), - ]; - for ar in aspect_ratios { - let request = make_request(Some("p"), ar, None, Some(2), None); - let distilled = distill_text(&request); - assert_eq!( - distilled.cost.cost_in_usd_cents, - Some(30), - "expected 30¢ regardless of aspect ratio (got {:?})", - ar, - ); - } - } - - #[test] - fn cost_is_independent_of_prompt() { - let with_prompt = distill_text(&make_request(Some("a cat"), None, None, Some(3), None)); - let without_prompt = distill_text(&make_request(None, None, None, Some(3), None)); - assert_eq!(with_prompt.cost.cost_in_usd_cents, Some(45)); - assert_eq!(without_prompt.cost.cost_in_usd_cents, Some(45)); - } - - #[test] - fn cost_metadata_flags_are_default() { - let request = make_request(Some("p"), None, None, Some(1), None); - let distilled = distill_text(&request); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - } - - mod edit { - use super::*; - - fn cost_for( - resolution: Option, - image_batch_count: Option, - num_image_refs: usize, - ) -> (Option, Option) { - let (tokens, hydration) = fake_image_refs(num_image_refs); - let request = make_request( - Some("make it pop"), - None, - resolution, - image_batch_count, - Some(tokens), - ); - let distilled = distill_edit(&request, &hydration); - ( - distilled.cost.cost_in_credits, - distilled.cost.cost_in_usd_cents, - ) - } - - // ── Default resolution (None → 1K → 15¢/image) ───────────────────────── - - #[test] - fn default_resolution_default_batch_costs_15_cents() { - let (credits, cents) = cost_for(None, None, 1); - assert_eq!(credits, Some(15)); - assert_eq!(cents, Some(15)); - } - - #[test] - fn default_resolution_batch_of_one_costs_15_cents() { - let (_c, cents) = cost_for(None, Some(1), 1); - assert_eq!(cents, Some(15)); - } - - #[test] - fn default_resolution_batch_of_two_costs_30_cents() { - let (_c, cents) = cost_for(None, Some(2), 1); - assert_eq!(cents, Some(30)); - } - - #[test] - fn default_resolution_batch_of_three_costs_45_cents() { - let (_c, cents) = cost_for(None, Some(3), 1); - assert_eq!(cents, Some(45)); - } - - #[test] - fn default_resolution_batch_of_four_costs_60_cents() { - let (_c, cents) = cost_for(None, Some(4), 1); - assert_eq!(cents, Some(60)); - } - - #[test] - fn default_resolution_batch_above_max_clamps_to_60_cents() { - let (_c, cents) = cost_for(None, Some(9), 1); - assert_eq!(cents, Some(60)); - } - - // ── 1K (15¢/image) ───────────────────────────────────────────────────── - - #[test] - fn one_k_batch_of_one_costs_15_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(1), 1); - assert_eq!(cents, Some(15)); - } - - #[test] - fn one_k_batch_of_four_costs_60_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::OneK), Some(4), 1); - assert_eq!(cents, Some(60)); - } - - // ── 2K (15¢/image) ───────────────────────────────────────────────────── - - #[test] - fn two_k_batch_of_one_costs_15_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(1), 1); - assert_eq!(cents, Some(15)); - } - - #[test] - fn two_k_batch_of_four_costs_60_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::TwoK), Some(4), 1); - assert_eq!(cents, Some(60)); - } - - // ── 4K (30¢/image) ───────────────────────────────────────────────────── - - #[test] - fn four_k_batch_of_one_costs_30_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(1), 1); - assert_eq!(cents, Some(30)); - } - - #[test] - fn four_k_batch_of_two_costs_60_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(2), 1); - assert_eq!(cents, Some(60)); - } - - #[test] - fn four_k_batch_of_four_costs_120_cents() { - let (_c, cents) = cost_for(Some(CommonResolution::FourK), Some(4), 1); - assert_eq!(cents, Some(120)); - } - - // ── Fallbacks ────────────────────────────────────────────────────────── - - #[test] - fn half_k_falls_back_to_one_k_pricing() { - let (_c, cents) = cost_for(Some(CommonResolution::HalfK), Some(1), 1); - assert_eq!(cents, Some(15)); - } - - #[test] - fn three_k_falls_back_to_two_k_pricing() { - let (_c, cents) = cost_for(Some(CommonResolution::ThreeK), Some(1), 1); - assert_eq!(cents, Some(15)); - } - - #[test] - fn cost_is_independent_of_image_ref_count() { - for num_refs in [1usize, 2, 3, 5] { - let (_c, cents) = cost_for(None, Some(2), num_refs); - assert_eq!( - cents, - Some(30), - "expected 30¢ regardless of {} image refs", - num_refs, - ); - } - } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let aspect_ratios = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::Auto), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::TallNineBySixteen), - ]; - for ar in aspect_ratios { - let (tokens, hydration) = fake_image_refs(2); - let request = make_request(Some("p"), ar, None, Some(2), Some(tokens)); - let distilled = distill_edit(&request, &hydration); - assert_eq!( - distilled.cost.cost_in_usd_cents, - Some(30), - "expected 30¢ regardless of aspect ratio (got {:?})", - ar, - ); - } - } - - #[test] - fn cost_metadata_flags_are_default() { - let (tokens, hydration) = fake_image_refs(1); - let request = make_request(Some("p"), None, None, Some(1), Some(tokens)); - let distilled = distill_edit(&request, &hydration); - assert!(!distilled.cost.is_free); - assert!(!distilled.cost.is_unlimited); - assert!(!distilled.cost.is_rate_limited); - assert!(!distilled.cost.has_watermark); - } - - #[test] - fn edit_and_text_cost_match_at_each_resolution() { - let resolutions = [ - None, - Some(CommonResolution::OneK), - Some(CommonResolution::TwoK), - Some(CommonResolution::FourK), - ]; - for res in resolutions { - for batch in [1u16, 2, 3, 4] { - let text = distill_text(&make_request(Some("p"), None, res, Some(batch), None)); - let (tokens, hydration) = fake_image_refs(1); - let edit = distill_edit( - &make_request(Some("p"), None, res, Some(batch), Some(tokens)), - &hydration, - ); - assert_eq!( - text.cost.cost_in_usd_cents, - edit.cost.cost_in_usd_cents, - "text/edit cost diverged at res={:?} batch={}", - res, - batch, - ); - } - } - } - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - // - // The distilled plan is `ImageGenerationPlan::FalNanaBananaPro(PlanFalNanaBananaPro)`. - // The handler picks t2i vs edit based on whether `image_urls` is empty. - // - // Pro's text-to-image endpoint has NO Auto aspect ratio — Auto* in t2i - // mode falls back to OneByOne. Edit mode does support Auto. - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_nano_banana_pro::{ - FalNbpNumImages, FalNbpResolution, PlanFalNanaBananaPro, - }; - use fal_client::requests::webhook::image::edit::enqueue_nano_banana_pro_edit_image_webhook::EnqueueNanoBananaProEditImageAspectRatio; - use fal_client::requests::webhook::image::text::enqueue_nano_banana_pro_text_to_image_webhook::EnqueueNanoBananaProTextToImageAspectRatio; - - fn with_text_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill_text(request); - match distilled.plan() { - ImageGenerationPlan::FalNanaBananaPro(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalNanaBananaPro, got {:?}", other), - } - } - - fn with_edit_plan( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - assertion: F, - ) { - let distilled = distill_edit(request, hydration); - match distilled.plan() { - ImageGenerationPlan::FalNanaBananaPro(plan) => assertion(plan), - other => panic!("expected ImageGenerationPlan::FalNanaBananaPro, got {:?}", other), - } - } - - mod text { - use super::*; - - fn assert_t2i_for_aspect_ratio( - ar: Option, - assertion: F, - ) { - with_text_plan(&make_request(Some("p"), ar, None, Some(1), None), assertion); - } - - fn assert_t2i_for_batch( - batch: Option, - assertion: F, - ) { - with_text_plan(&make_request(Some("p"), None, None, batch, None), assertion); - } - - fn assert_t2i_for_resolution( - res: Option, - assertion: F, - ) { - with_text_plan(&make_request(Some("p"), None, res, Some(1), None), assertion); - } - - // ── Mode detection ──────────────────────────────────────────────────── - - #[test] - fn text_mode_has_empty_image_urls() { - with_text_plan(&make_request(Some("p"), None, None, Some(1), None), |plan| { - assert!(plan.image_urls.is_empty(), "text mode must have no image_urls"); - }); - } - - // ── Aspect ratio direct mappings ────────────────────────────────────── - - #[test] - fn default_aspect_ratio_is_none() { - assert_t2i_for_aspect_ratio(None, |plan| { - assert!(plan.t2i_aspect_ratio.is_none()); - }); - } - - #[test] - fn square_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn square_hd_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn wide_5x4_yields_five_by_four() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::FiveByFour) - )); - }); - } - - #[test] - fn wide_4x3_yields_four_by_three() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::FourByThree) - )); - }); - } - - #[test] - fn wide_3x2_yields_three_by_two() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideThreeByTwo), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::ThreeByTwo) - )); - }); - } - - #[test] - fn wide_16x9_yields_sixteen_by_nine() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_alias_yields_sixteen_by_nine() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_21x9_yields_twenty_one_by_nine() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::WideTwentyOneByNine), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::TwentyOneByNine) - )); - }); - } - - #[test] - fn tall_4x5_yields_four_by_five() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallFourByFive), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::FourByFive) - )); - }); - } - - #[test] - fn tall_3x4_yields_three_by_four() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::ThreeByFour) - )); - }); - } - - #[test] - fn tall_2x3_yields_two_by_three() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::TwoByThree) - )); - }); - } - - #[test] - fn tall_9x16_yields_nine_by_sixteen() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::NineBySixteen) - )); - }); - } - - #[test] - fn tall_alias_yields_nine_by_sixteen() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::NineBySixteen) - )); - }); - } - - // ── Auto* falls back to OneByOne in t2i (Pro has no t2i Auto) ───────── - - #[test] - fn auto_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn auto_2k_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Auto2k), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn auto_4k_yields_one_by_one() { - assert_t2i_for_aspect_ratio(Some(CommonAspectRatio::Auto4k), |plan| { - assert!(matches!( - plan.t2i_aspect_ratio, - Some(EnqueueNanoBananaProTextToImageAspectRatio::OneByOne) - )); - }); - } - - // ── Resolution mapping ──────────────────────────────────────────────── - - #[test] - fn default_resolution_is_none() { - assert_t2i_for_resolution(None, |plan| { - assert!(plan.resolution.is_none()); - }); - } - - #[test] - fn one_k_yields_one_k() { - assert_t2i_for_resolution(Some(CommonResolution::OneK), |plan| { - assert!(matches!(plan.resolution, Some(FalNbpResolution::OneK))); - }); - } - - #[test] - fn two_k_yields_two_k() { - assert_t2i_for_resolution(Some(CommonResolution::TwoK), |plan| { - assert!(matches!(plan.resolution, Some(FalNbpResolution::TwoK))); - }); - } - - #[test] - fn four_k_yields_four_k() { - assert_t2i_for_resolution(Some(CommonResolution::FourK), |plan| { - assert!(matches!(plan.resolution, Some(FalNbpResolution::FourK))); - }); - } - - #[test] - fn three_k_falls_back_to_two_k() { - assert_t2i_for_resolution(Some(CommonResolution::ThreeK), |plan| { - assert!(matches!(plan.resolution, Some(FalNbpResolution::TwoK))); - }); - } - - #[test] - fn half_k_falls_back_to_one_k() { - assert_t2i_for_resolution(Some(CommonResolution::HalfK), |plan| { - assert!(matches!(plan.resolution, Some(FalNbpResolution::OneK))); - }); - } - - // ── Num images mapping ──────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - assert_t2i_for_batch(None, |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::One)); - }); - } - - #[test] - fn batch_of_one_yields_one() { - assert_t2i_for_batch(Some(1), |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::One)); - }); - } - - #[test] - fn batch_of_two_yields_two() { - assert_t2i_for_batch(Some(2), |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::Two)); - }); - } - - #[test] - fn batch_of_three_yields_three() { - assert_t2i_for_batch(Some(3), |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::Three)); - }); - } - - #[test] - fn batch_of_four_yields_four() { - assert_t2i_for_batch(Some(4), |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::Four)); - }); - } - - #[test] - fn batch_above_four_clamps_to_four() { - assert_t2i_for_batch(Some(9), |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_text_plan(&make_request(Some("a corgi in a hat"), None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_text_plan(&make_request(None, None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } - - mod edit { - use super::*; - - fn assert_edit_for_aspect_ratio( - ar: Option, - assertion: F, - ) { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), ar, None, Some(1), Some(tokens)), - &hydration, - assertion, - ); - } - - fn assert_edit_for_batch( - batch: Option, - assertion: F, - ) { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, batch, Some(tokens)), - &hydration, - assertion, - ); - } - - fn assert_edit_for_resolution( - res: Option, - assertion: F, - ) { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, res, Some(1), Some(tokens)), - &hydration, - assertion, - ); - } - - // ── Mode detection / image-url passthrough ──────────────────────────── - - #[test] - fn edit_mode_populates_image_urls() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 1); - assert!(plan.image_urls[0].starts_with("https://fake.example.com/")); - }, - ); - } - - #[test] - fn edit_mode_with_two_image_refs() { - let (tokens, hydration) = fake_image_refs(2); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 2); - }, - ); - } - - #[test] - fn edit_mode_with_five_image_refs() { - let (tokens, hydration) = fake_image_refs(5); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 5); - }, - ); - } - - #[test] - fn edit_image_urls_are_hydrated_from_map() { - let (tokens, hydration) = fake_image_refs(3); - let expected: Vec = tokens - .iter() - .map(|t| hydration.get(t).unwrap().to_string()) - .collect(); - with_edit_plan( - &make_request(Some("p"), None, None, Some(1), Some(tokens.clone())), - &hydration, - |plan| { - assert_eq!(plan.image_urls, expected); - }, - ); - } - - // ── Aspect ratio direct mappings ────────────────────────────────────── - - #[test] - fn default_aspect_ratio_is_none() { - assert_edit_for_aspect_ratio(None, |plan| { - assert!(plan.edit_aspect_ratio.is_none()); - }); - } - - #[test] - fn square_yields_one_by_one() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Square), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn square_hd_yields_one_by_one() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::SquareHd), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::OneByOne) - )); - }); - } - - #[test] - fn wide_5x4_yields_five_by_four() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideFiveByFour), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::FiveByFour) - )); - }); - } - - #[test] - fn wide_4x3_yields_four_by_three() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideFourByThree), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::FourByThree) - )); - }); - } - - #[test] - fn wide_3x2_yields_three_by_two() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideThreeByTwo), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::ThreeByTwo) - )); - }); - } - - #[test] - fn wide_16x9_yields_sixteen_by_nine() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideSixteenByNine), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_alias_yields_sixteen_by_nine() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Wide), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::SixteenByNine) - )); - }); - } - - #[test] - fn wide_21x9_yields_twenty_one_by_nine() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::WideTwentyOneByNine), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::TwentyOneByNine) - )); - }); - } - - #[test] - fn tall_4x5_yields_four_by_five() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallFourByFive), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::FourByFive) - )); - }); - } - - #[test] - fn tall_3x4_yields_three_by_four() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallThreeByFour), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::ThreeByFour) - )); - }); - } - - #[test] - fn tall_2x3_yields_two_by_three() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallTwoByThree), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::TwoByThree) - )); - }); - } - - #[test] - fn tall_9x16_yields_nine_by_sixteen() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::TallNineBySixteen), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::NineBySixteen) - )); - }); - } - - #[test] - fn tall_alias_yields_nine_by_sixteen() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Tall), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::NineBySixteen) - )); - }); - } - - // ── Auto* maps to Auto in edit mode ─────────────────────────────────── - - #[test] - fn auto_yields_auto() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Auto), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::Auto) - )); - }); - } - - #[test] - fn auto_2k_yields_auto() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Auto2k), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::Auto) - )); - }); - } - - #[test] - fn auto_4k_yields_auto() { - assert_edit_for_aspect_ratio(Some(CommonAspectRatio::Auto4k), |plan| { - assert!(matches!( - plan.edit_aspect_ratio, - Some(EnqueueNanoBananaProEditImageAspectRatio::Auto) - )); - }); - } - - // ── Resolution mapping ──────────────────────────────────────────────── - - #[test] - fn default_resolution_is_none() { - assert_edit_for_resolution(None, |plan| { - assert!(plan.resolution.is_none()); - }); - } - - #[test] - fn one_k_yields_one_k() { - assert_edit_for_resolution(Some(CommonResolution::OneK), |plan| { - assert!(matches!(plan.resolution, Some(FalNbpResolution::OneK))); - }); - } - - #[test] - fn two_k_yields_two_k() { - assert_edit_for_resolution(Some(CommonResolution::TwoK), |plan| { - assert!(matches!(plan.resolution, Some(FalNbpResolution::TwoK))); - }); - } - - #[test] - fn four_k_yields_four_k() { - assert_edit_for_resolution(Some(CommonResolution::FourK), |plan| { - assert!(matches!(plan.resolution, Some(FalNbpResolution::FourK))); - }); - } - - #[test] - fn three_k_falls_back_to_two_k() { - assert_edit_for_resolution(Some(CommonResolution::ThreeK), |plan| { - assert!(matches!(plan.resolution, Some(FalNbpResolution::TwoK))); - }); - } - - #[test] - fn half_k_falls_back_to_one_k() { - assert_edit_for_resolution(Some(CommonResolution::HalfK), |plan| { - assert!(matches!(plan.resolution, Some(FalNbpResolution::OneK))); - }); - } - - // ── Num images mapping ──────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - assert_edit_for_batch(None, |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::One)); - }); - } - - #[test] - fn batch_of_one_yields_one() { - assert_edit_for_batch(Some(1), |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::One)); - }); - } - - #[test] - fn batch_of_two_yields_two() { - assert_edit_for_batch(Some(2), |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::Two)); - }); - } - - #[test] - fn batch_of_three_yields_three() { - assert_edit_for_batch(Some(3), |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::Three)); - }); - } - - #[test] - fn batch_of_four_yields_four() { - assert_edit_for_batch(Some(4), |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::Four)); - }); - } - - #[test] - fn batch_above_four_clamps_to_four() { - assert_edit_for_batch(Some(9), |plan| { - assert!(matches!(plan.num_images, FalNbpNumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("make it shiny"), None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.prompt, Some("make it shiny".to_string())); - }, - ); - } - - #[test] - fn missing_prompt_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(None, None, None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.prompt, None); - }, - ); - } - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4.rs deleted file mode 100644 index 6f5b7f056a..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4.rs +++ /dev/null @@ -1,613 +0,0 @@ -//! Tests for omni-gen against the Seedream 4 model. -//! -//! Seedream 4 is multi-function: text-to-image and image-to-image (edit). -//! -//! Artcraft-tier pricing: flat 3¢/image regardless of size, quality, or mode. -//! -//! The Fal execution plan maps aspect ratios to Seedream-native image sizes: -//! Square, SquareHd, LandscapeFourThree, LandscapeSixteenNine, -//! PortraitFourThree, PortraitSixteenNine, Auto, Auto2k, Auto4k. -//! Auto in t2i mode falls back to Square; in edit mode it passes through. -//! Unsupported ratios (5:4, 3:2, 21:9, 4:5, 2:3, 9:21) fall back under -//! PayMoreUpgrade to the nearest 16:9 or 9:16 variant. - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - image_batch_count: Option, - image_media_tokens: Option>, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::Seedream4), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens, - resolution: None, - aspect_ratio, - quality: None, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - fn fake_image_refs(n: usize) -> (Vec, HashMap) { - let mut tokens = Vec::with_capacity(n); - let mut map = HashMap::with_capacity(n); - for i in 0..n { - let token = MediaFileToken::new_from_str(&format!("mf_test{:028}", i)); - let url = Url::parse(&format!("https://fake.example.com/img{}.png", i)).unwrap(); - map.insert(token.clone(), url); - tokens.push(token); - } - (tokens, map) - } - - fn distill_text(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for Seedream 4 (text)") - } - - fn distill_edit( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - ) -> DistilledImageRequest { - distill_image_request(request, Some(hydration)) - .expect("distill_image_request should succeed for Seedream 4 (edit)") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST — flat 3¢/image - // ──────────────────────────────────────────────────────────────────────────── - mod cost { - use super::*; - - mod text { - use super::*; - - fn cost_for_batch(batch: Option) -> u64 { - let request = make_request(Some("p"), None, batch, None); - distill_text(&request).cost.cost_in_usd_cents.unwrap() - } - - #[test] - fn default_batch_costs_3_cents() { - assert_eq!(cost_for_batch(None), 3); - } - - #[test] - fn batch_of_one_costs_3() { assert_eq!(cost_for_batch(Some(1)), 3); } - - #[test] - fn batch_of_two_costs_6() { assert_eq!(cost_for_batch(Some(2)), 6); } - - #[test] - fn batch_of_three_costs_9() { assert_eq!(cost_for_batch(Some(3)), 9); } - - #[test] - fn batch_of_four_costs_12() { assert_eq!(cost_for_batch(Some(4)), 12); } - - #[test] - fn batch_above_max_clamps_to_12() { assert_eq!(cost_for_batch(Some(7)), 12); } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let ars = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::SquareHd), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::WideFourByThree), - Some(CommonAspectRatio::TallNineBySixteen), - Some(CommonAspectRatio::TallThreeByFour), - Some(CommonAspectRatio::Auto), - Some(CommonAspectRatio::Auto2k), - Some(CommonAspectRatio::Auto4k), - ]; - for ar in ars { - let request = make_request(Some("p"), ar, Some(2), None); - assert_eq!( - distill_text(&request).cost.cost_in_usd_cents.unwrap(), 6, - "expected 6¢ regardless of aspect ratio {:?}", ar, - ); - } - } - - #[test] - fn cost_is_independent_of_prompt() { - let a = distill_text(&make_request(Some("a cat"), None, Some(1), None)); - let b = distill_text(&make_request(None, None, Some(1), None)); - assert_eq!(a.cost.cost_in_usd_cents, b.cost.cost_in_usd_cents); - } - - #[test] - fn cost_metadata_flags_are_default() { - let d = distill_text(&make_request(Some("p"), None, Some(1), None)); - assert!(!d.cost.is_free); - assert!(!d.cost.is_unlimited); - assert!(!d.cost.is_rate_limited); - assert!(!d.cost.has_watermark); - } - } - - mod edit { - use super::*; - - fn cost_edit(batch: Option, num_refs: usize) -> u64 { - let (tokens, hydration) = fake_image_refs(num_refs); - let request = make_request(Some("p"), None, batch, Some(tokens)); - distill_edit(&request, &hydration).cost.cost_in_usd_cents.unwrap() - } - - #[test] - fn default_batch_costs_3() { assert_eq!(cost_edit(None, 1), 3); } - - #[test] - fn batch_of_one_costs_3() { assert_eq!(cost_edit(Some(1), 1), 3); } - - #[test] - fn batch_of_two_costs_6() { assert_eq!(cost_edit(Some(2), 1), 6); } - - #[test] - fn batch_of_four_costs_12() { assert_eq!(cost_edit(Some(4), 1), 12); } - - #[test] - fn batch_above_max_clamps_to_12() { assert_eq!(cost_edit(Some(9), 1), 12); } - - #[test] - fn cost_independent_of_input_image_count() { - for num_refs in [1usize, 2, 3, 5] { - assert_eq!(cost_edit(Some(2), num_refs), 6, "expected 6¢ regardless of {} input images", num_refs); - } - } - - #[test] - fn edit_and_text_cost_match() { - for batch in [1u16, 2, 3, 4] { - let text = distill_text(&make_request(Some("p"), None, Some(batch), None)) - .cost.cost_in_usd_cents.unwrap(); - let (tokens, hydration) = fake_image_refs(1); - let edit = distill_edit( - &make_request(Some("p"), None, Some(batch), Some(tokens)), - &hydration, - ).cost.cost_in_usd_cents.unwrap(); - assert_eq!(text, edit, "text/edit cost diverged at batch={}", batch); - } - } - - #[test] - fn cost_metadata_flags_are_default() { - let (tokens, hydration) = fake_image_refs(1); - let d = distill_edit(&make_request(Some("p"), None, Some(1), Some(tokens)), &hydration); - assert!(!d.cost.is_free); - assert!(!d.cost.is_unlimited); - assert!(!d.cost.is_rate_limited); - assert!(!d.cost.has_watermark); - } - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_seedream_4::{ - FalSeedream4ImageSize, FalSeedream4NumImages, PlanFalSeedream4, - }; - - fn with_text_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill_text(request); - match distilled.plan() { - ImageGenerationPlan::FalSeedream4(plan) => assertion(plan), - other => panic!("expected FalSeedream4, got {:?}", other), - } - } - - fn with_edit_plan( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - assertion: F, - ) { - let distilled = distill_edit(request, hydration); - match distilled.plan() { - ImageGenerationPlan::FalSeedream4(plan) => assertion(plan), - other => panic!("expected FalSeedream4, got {:?}", other), - } - } - - mod text { - use super::*; - - // ── Mode detection ──────────────────────────────────────────────────── - - #[test] - fn text_mode_has_empty_image_urls() { - with_text_plan(&make_request(Some("p"), None, Some(1), None), |plan| { - assert!(plan.image_urls.is_empty()); - }); - } - - // ── Image size: direct mappings ─────────────────────────────────────── - - #[test] - fn default_image_size_is_none() { - with_text_plan(&make_request(Some("p"), None, Some(1), None), |plan| { - assert!(plan.image_size.is_none()); - }); - } - - #[test] - fn square_yields_square() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Square), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::Square))); - }); - } - - #[test] - fn square_hd_yields_square_hd() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::SquareHd), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::SquareHd))); - }); - } - - #[test] - fn wide_16x9_yields_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideSixteenByNine), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_alias_yields_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Wide), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_4x3_yields_landscape_four_three() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideFourByThree), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::LandscapeFourThree))); - }); - } - - #[test] - fn tall_9x16_yields_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallNineBySixteen), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_alias_yields_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Tall), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_3x4_yields_portrait_four_three() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallThreeByFour), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::PortraitFourThree))); - }); - } - - // ── Image size: Auto variants ───────────────────────────────────────── - - #[test] - fn auto_in_text_mode_yields_square() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Auto), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::Square))); - }); - } - - #[test] - fn auto_2k_yields_auto_2k() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Auto2k), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::Auto2k))); - }); - } - - #[test] - fn auto_4k_yields_auto_4k() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Auto4k), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::Auto4k))); - }); - } - - // ── Image size: fallbacks (unsupported → nearest under PayMoreUpgrade) - - #[test] - fn wide_5x4_falls_back_to_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideFiveByFour), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_3x2_falls_back_to_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideThreeByTwo), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_21x9_falls_back_to_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideTwentyOneByNine), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn tall_4x5_falls_back_to_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallFourByFive), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_2x3_falls_back_to_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallTwoByThree), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_9x21_falls_back_to_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallNineByTwentyOne), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::PortraitSixteenNine))); - }); - } - - // ── Num images ──────────────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - with_text_plan(&make_request(Some("p"), None, None, None), |plan| { - assert!(matches!(plan.num_images, FalSeedream4NumImages::One)); - }); - } - - #[test] - fn batch_direct_mapping() { - for (count, expected) in [ - (1u16, FalSeedream4NumImages::One), (2, FalSeedream4NumImages::Two), - (3, FalSeedream4NumImages::Three), (4, FalSeedream4NumImages::Four), - ] { - with_text_plan(&make_request(Some("p"), None, Some(count), None), |plan| { - assert!(std::mem::discriminant(&plan.num_images) == std::mem::discriminant(&expected)); - }); - } - } - - #[test] - fn batch_above_four_clamps_to_four() { - with_text_plan(&make_request(Some("p"), None, Some(9), None), |plan| { - assert!(matches!(plan.num_images, FalSeedream4NumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_text_plan(&make_request(Some("a corgi in a hat"), None, Some(1), None), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_text_plan(&make_request(None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } - - mod edit { - use super::*; - - // ── Mode detection / image-url passthrough ──────────────────────────── - - #[test] - fn edit_mode_populates_image_urls() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 1); - assert!(plan.image_urls[0].starts_with("https://fake.example.com/")); - }, - ); - } - - #[test] - fn edit_mode_with_three_image_refs() { - let (tokens, hydration) = fake_image_refs(3); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.image_urls.len(), 3); }, - ); - } - - #[test] - fn edit_mode_with_five_image_refs() { - let (tokens, hydration) = fake_image_refs(5); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.image_urls.len(), 5); }, - ); - } - - #[test] - fn edit_image_urls_are_hydrated_from_map() { - let (tokens, hydration) = fake_image_refs(3); - let expected: Vec = tokens.iter() - .map(|t| hydration.get(t).unwrap().to_string()) - .collect(); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens.clone())), - &hydration, - |plan| { assert_eq!(plan.image_urls, expected); }, - ); - } - - // ── Auto in edit mode yields Auto (not Square) ──────────────────────── - - #[test] - fn auto_in_edit_mode_yields_auto() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Auto), Some(1), Some(tokens)), - &hydration, - |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::Auto))); - }, - ); - } - - // ── Image size direct mappings in edit mode ─────────────────────────── - - #[test] - fn edit_default_image_size_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(plan.image_size.is_none()); }, - ); - } - - #[test] - fn edit_square_yields_square() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Square), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::Square))); }, - ); - } - - #[test] - fn edit_wide_16x9_yields_landscape_sixteen_nine() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::WideSixteenByNine), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::LandscapeSixteenNine))); }, - ); - } - - #[test] - fn edit_tall_9x16_yields_portrait_sixteen_nine() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::TallNineBySixteen), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::PortraitSixteenNine))); }, - ); - } - - #[test] - fn edit_auto_2k() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Auto2k), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::Auto2k))); }, - ); - } - - #[test] - fn edit_auto_4k() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Auto4k), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream4ImageSize::Auto4k))); }, - ); - } - - // ── Num images in edit mode ─────────────────────────────────────────── - - #[test] - fn edit_default_batch_count_is_one() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalSeedream4NumImages::One)); }, - ); - } - - #[test] - fn edit_batch_of_four_yields_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(4), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalSeedream4NumImages::Four)); }, - ); - } - - #[test] - fn edit_batch_above_four_clamps_to_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(9), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalSeedream4NumImages::Four)); }, - ); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn edit_prompt_is_passed_through() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("make it shiny"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, Some("make it shiny".to_string())); }, - ); - } - - #[test] - fn edit_missing_prompt_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, None); }, - ); - } - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4p5.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4p5.rs deleted file mode 100644 index d03f0a4a0b..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4p5.rs +++ /dev/null @@ -1,610 +0,0 @@ -//! Tests for omni-gen against the Seedream 4.5 model. -//! -//! Seedream 4.5 is multi-function: text-to-image and image-to-image (edit). -//! -//! Artcraft-tier pricing: flat 4¢/image regardless of size, quality, or mode. -//! -//! The Fal execution plan maps aspect ratios to Seedream-native image sizes. -//! Unlike V4, V4.5 has no bare "Auto" — both `Auto` and `Auto2k` map to -//! `Auto2k`, and `Auto4k` maps to `Auto4k`. -//! Unsupported ratios (5:4, 3:2, 21:9, 4:5, 2:3, 9:21) fall back under -//! PayMoreUpgrade to the nearest 16:9 or 9:16 variant. - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - image_batch_count: Option, - image_media_tokens: Option>, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::Seedream4p5), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens, - resolution: None, - aspect_ratio, - quality: None, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - fn fake_image_refs(n: usize) -> (Vec, HashMap) { - let mut tokens = Vec::with_capacity(n); - let mut map = HashMap::with_capacity(n); - for i in 0..n { - let token = MediaFileToken::new_from_str(&format!("mf_test{:028}", i)); - let url = Url::parse(&format!("https://fake.example.com/img{}.png", i)).unwrap(); - map.insert(token.clone(), url); - tokens.push(token); - } - (tokens, map) - } - - fn distill_text(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for Seedream 4.5 (text)") - } - - fn distill_edit( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - ) -> DistilledImageRequest { - distill_image_request(request, Some(hydration)) - .expect("distill_image_request should succeed for Seedream 4.5 (edit)") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST — flat 4¢/image - // ──────────────────────────────────────────────────────────────────────────── - mod cost { - use super::*; - - mod text { - use super::*; - - fn cost_for_batch(batch: Option) -> u64 { - let request = make_request(Some("p"), None, batch, None); - distill_text(&request).cost.cost_in_usd_cents.unwrap() - } - - #[test] - fn default_batch_costs_4_cents() { assert_eq!(cost_for_batch(None), 4); } - - #[test] - fn batch_of_one_costs_4() { assert_eq!(cost_for_batch(Some(1)), 4); } - - #[test] - fn batch_of_two_costs_8() { assert_eq!(cost_for_batch(Some(2)), 8); } - - #[test] - fn batch_of_three_costs_12() { assert_eq!(cost_for_batch(Some(3)), 12); } - - #[test] - fn batch_of_four_costs_16() { assert_eq!(cost_for_batch(Some(4)), 16); } - - #[test] - fn batch_above_max_clamps_to_16() { assert_eq!(cost_for_batch(Some(7)), 16); } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let ars = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::SquareHd), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::WideFourByThree), - Some(CommonAspectRatio::TallNineBySixteen), - Some(CommonAspectRatio::TallThreeByFour), - Some(CommonAspectRatio::Auto), - Some(CommonAspectRatio::Auto2k), - Some(CommonAspectRatio::Auto4k), - ]; - for ar in ars { - let request = make_request(Some("p"), ar, Some(2), None); - assert_eq!( - distill_text(&request).cost.cost_in_usd_cents.unwrap(), 8, - "expected 8¢ regardless of aspect ratio {:?}", ar, - ); - } - } - - #[test] - fn cost_is_independent_of_prompt() { - let a = distill_text(&make_request(Some("a cat"), None, Some(1), None)); - let b = distill_text(&make_request(None, None, Some(1), None)); - assert_eq!(a.cost.cost_in_usd_cents, b.cost.cost_in_usd_cents); - } - - #[test] - fn cost_metadata_flags_are_default() { - let d = distill_text(&make_request(Some("p"), None, Some(1), None)); - assert!(!d.cost.is_free); - assert!(!d.cost.is_unlimited); - assert!(!d.cost.is_rate_limited); - assert!(!d.cost.has_watermark); - } - } - - mod edit { - use super::*; - - fn cost_edit(batch: Option, num_refs: usize) -> u64 { - let (tokens, hydration) = fake_image_refs(num_refs); - let request = make_request(Some("p"), None, batch, Some(tokens)); - distill_edit(&request, &hydration).cost.cost_in_usd_cents.unwrap() - } - - #[test] - fn default_batch_costs_4() { assert_eq!(cost_edit(None, 1), 4); } - - #[test] - fn batch_of_one_costs_4() { assert_eq!(cost_edit(Some(1), 1), 4); } - - #[test] - fn batch_of_two_costs_8() { assert_eq!(cost_edit(Some(2), 1), 8); } - - #[test] - fn batch_of_four_costs_16() { assert_eq!(cost_edit(Some(4), 1), 16); } - - #[test] - fn batch_above_max_clamps_to_16() { assert_eq!(cost_edit(Some(9), 1), 16); } - - #[test] - fn cost_independent_of_input_image_count() { - for num_refs in [1usize, 2, 3, 5] { - assert_eq!(cost_edit(Some(2), num_refs), 8, "expected 8¢ regardless of {} input images", num_refs); - } - } - - #[test] - fn edit_and_text_cost_match() { - for batch in [1u16, 2, 3, 4] { - let text = distill_text(&make_request(Some("p"), None, Some(batch), None)) - .cost.cost_in_usd_cents.unwrap(); - let (tokens, hydration) = fake_image_refs(1); - let edit = distill_edit( - &make_request(Some("p"), None, Some(batch), Some(tokens)), - &hydration, - ).cost.cost_in_usd_cents.unwrap(); - assert_eq!(text, edit, "text/edit cost diverged at batch={}", batch); - } - } - - #[test] - fn cost_metadata_flags_are_default() { - let (tokens, hydration) = fake_image_refs(1); - let d = distill_edit(&make_request(Some("p"), None, Some(1), Some(tokens)), &hydration); - assert!(!d.cost.is_free); - assert!(!d.cost.is_unlimited); - assert!(!d.cost.is_rate_limited); - assert!(!d.cost.has_watermark); - } - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_seedream_4p5::{ - FalSeedream4p5ImageSize, FalSeedream4p5NumImages, PlanFalSeedream4p5, - }; - - fn with_text_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill_text(request); - match distilled.plan() { - ImageGenerationPlan::FalSeedream4p5(plan) => assertion(plan), - other => panic!("expected FalSeedream4p5, got {:?}", other), - } - } - - fn with_edit_plan( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - assertion: F, - ) { - let distilled = distill_edit(request, hydration); - match distilled.plan() { - ImageGenerationPlan::FalSeedream4p5(plan) => assertion(plan), - other => panic!("expected FalSeedream4p5, got {:?}", other), - } - } - - mod text { - use super::*; - - // ── Mode detection ──────────────────────────────────────────────────── - - #[test] - fn text_mode_has_empty_image_urls() { - with_text_plan(&make_request(Some("p"), None, Some(1), None), |plan| { - assert!(plan.image_urls.is_empty()); - }); - } - - // ── Image size: direct mappings ─────────────────────────────────────── - - #[test] - fn default_image_size_is_none() { - with_text_plan(&make_request(Some("p"), None, Some(1), None), |plan| { - assert!(plan.image_size.is_none()); - }); - } - - #[test] - fn square_yields_square() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Square), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::Square))); - }); - } - - #[test] - fn square_hd_yields_square_hd() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::SquareHd), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::SquareHd))); - }); - } - - #[test] - fn wide_16x9_yields_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideSixteenByNine), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_alias_yields_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Wide), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_4x3_yields_landscape_four_three() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideFourByThree), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::LandscapeFourThree))); - }); - } - - #[test] - fn tall_9x16_yields_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallNineBySixteen), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_alias_yields_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Tall), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_3x4_yields_portrait_four_three() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallThreeByFour), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::PortraitFourThree))); - }); - } - - // ── Image size: Auto variants (no bare Auto — maps to Auto2k) ───────── - - #[test] - fn auto_yields_auto_2k() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Auto), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::Auto2k))); - }); - } - - #[test] - fn auto_2k_yields_auto_2k() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Auto2k), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::Auto2k))); - }); - } - - #[test] - fn auto_4k_yields_auto_4k() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Auto4k), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::Auto4k))); - }); - } - - // ── Image size: fallbacks ───────────────────────────────────────────── - - #[test] - fn wide_5x4_falls_back_to_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideFiveByFour), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_3x2_falls_back_to_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideThreeByTwo), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_21x9_falls_back_to_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideTwentyOneByNine), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn tall_4x5_falls_back_to_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallFourByFive), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_2x3_falls_back_to_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallTwoByThree), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_9x21_falls_back_to_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallNineByTwentyOne), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::PortraitSixteenNine))); - }); - } - - // ── Num images ──────────────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - with_text_plan(&make_request(Some("p"), None, None, None), |plan| { - assert!(matches!(plan.num_images, FalSeedream4p5NumImages::One)); - }); - } - - #[test] - fn batch_direct_mapping() { - for (count, expected) in [ - (1u16, FalSeedream4p5NumImages::One), (2, FalSeedream4p5NumImages::Two), - (3, FalSeedream4p5NumImages::Three), (4, FalSeedream4p5NumImages::Four), - ] { - with_text_plan(&make_request(Some("p"), None, Some(count), None), |plan| { - assert!(std::mem::discriminant(&plan.num_images) == std::mem::discriminant(&expected)); - }); - } - } - - #[test] - fn batch_above_four_clamps_to_four() { - with_text_plan(&make_request(Some("p"), None, Some(9), None), |plan| { - assert!(matches!(plan.num_images, FalSeedream4p5NumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_text_plan(&make_request(Some("a corgi in a hat"), None, Some(1), None), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_text_plan(&make_request(None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } - - mod edit { - use super::*; - - // ── Mode detection / image-url passthrough ──────────────────────────── - - #[test] - fn edit_mode_populates_image_urls() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 1); - assert!(plan.image_urls[0].starts_with("https://fake.example.com/")); - }, - ); - } - - #[test] - fn edit_mode_with_three_image_refs() { - let (tokens, hydration) = fake_image_refs(3); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.image_urls.len(), 3); }, - ); - } - - #[test] - fn edit_mode_with_five_image_refs() { - let (tokens, hydration) = fake_image_refs(5); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.image_urls.len(), 5); }, - ); - } - - #[test] - fn edit_image_urls_are_hydrated_from_map() { - let (tokens, hydration) = fake_image_refs(3); - let expected: Vec = tokens.iter() - .map(|t| hydration.get(t).unwrap().to_string()) - .collect(); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens.clone())), - &hydration, - |plan| { assert_eq!(plan.image_urls, expected); }, - ); - } - - // ── Auto in edit mode also maps to Auto2k (no bare Auto in v4.5) ────── - - #[test] - fn auto_in_edit_mode_yields_auto_2k() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Auto), Some(1), Some(tokens)), - &hydration, - |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::Auto2k))); - }, - ); - } - - // ── Image size direct mappings in edit mode ─────────────────────────── - - #[test] - fn edit_default_image_size_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(plan.image_size.is_none()); }, - ); - } - - #[test] - fn edit_square_yields_square() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Square), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::Square))); }, - ); - } - - #[test] - fn edit_wide_16x9_yields_landscape_sixteen_nine() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::WideSixteenByNine), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::LandscapeSixteenNine))); }, - ); - } - - #[test] - fn edit_tall_9x16_yields_portrait_sixteen_nine() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::TallNineBySixteen), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::PortraitSixteenNine))); }, - ); - } - - #[test] - fn edit_auto_2k() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Auto2k), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::Auto2k))); }, - ); - } - - #[test] - fn edit_auto_4k() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Auto4k), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream4p5ImageSize::Auto4k))); }, - ); - } - - // ── Num images in edit mode ─────────────────────────────────────────── - - #[test] - fn edit_default_batch_count_is_one() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalSeedream4p5NumImages::One)); }, - ); - } - - #[test] - fn edit_batch_of_four_yields_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(4), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalSeedream4p5NumImages::Four)); }, - ); - } - - #[test] - fn edit_batch_above_four_clamps_to_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(9), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalSeedream4p5NumImages::Four)); }, - ); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn edit_prompt_is_passed_through() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("make it shiny"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, Some("make it shiny".to_string())); }, - ); - } - - #[test] - fn edit_missing_prompt_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, None); }, - ); - } - } - } -} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_5_lite.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_5_lite.rs deleted file mode 100644 index d92ec7bb03..0000000000 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_5_lite.rs +++ /dev/null @@ -1,610 +0,0 @@ -//! Tests for omni-gen against the Seedream 5 Lite model. -//! -//! Seedream 5 Lite is multi-function: text-to-image and image-to-image (edit). -//! -//! Artcraft-tier pricing: flat 4¢/image regardless of size, quality, or mode. -//! -//! The Fal execution plan maps aspect ratios to Seedream-native image sizes. -//! V5 Lite supports Auto2k and Auto3k (no bare Auto, no Auto4k): -//! - Auto / Auto2k → Auto2k -//! - Auto4k → Auto3k (fallback) -//! Unsupported ratios (5:4, 3:2, 21:9, 4:5, 2:3, 9:21) fall back under -//! PayMoreUpgrade to the nearest 16:9 or 9:16 variant. - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - - use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; - use enums::common::generation::common_aspect_ratio::CommonAspectRatio; - use enums::common::generation::common_image_model::CommonImageModel; - use tokens::tokens::media_files::MediaFileToken; - use url::Url; - - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ - distill_image_request, DistilledImageRequest, - }; - - // ── Helpers ──────────────────────────────────────────────────────────────── - - fn make_request( - prompt: Option<&str>, - aspect_ratio: Option, - image_batch_count: Option, - image_media_tokens: Option>, - ) -> OmniGenImageCostAndGenerateRequest { - OmniGenImageCostAndGenerateRequest { - idempotency_token: Some("11111111-1111-1111-1111-111111111111".to_string()), - model: Some(CommonImageModel::Seedream5Lite), - prompt: prompt.map(|s| s.to_string()), - image_media_tokens, - resolution: None, - aspect_ratio, - quality: None, - image_batch_count, - adjust_horizontal_angle: None, - adjust_vertical_angle: None, - adjust_zoom: None, - } - } - - fn fake_image_refs(n: usize) -> (Vec, HashMap) { - let mut tokens = Vec::with_capacity(n); - let mut map = HashMap::with_capacity(n); - for i in 0..n { - let token = MediaFileToken::new_from_str(&format!("mf_test{:028}", i)); - let url = Url::parse(&format!("https://fake.example.com/img{}.png", i)).unwrap(); - map.insert(token.clone(), url); - tokens.push(token); - } - (tokens, map) - } - - fn distill_text(request: &OmniGenImageCostAndGenerateRequest) -> DistilledImageRequest { - let empty: HashMap = HashMap::new(); - distill_image_request(request, Some(&empty)) - .expect("distill_image_request should succeed for Seedream 5 Lite (text)") - } - - fn distill_edit( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - ) -> DistilledImageRequest { - distill_image_request(request, Some(hydration)) - .expect("distill_image_request should succeed for Seedream 5 Lite (edit)") - } - - // ──────────────────────────────────────────────────────────────────────────── - // COST — flat 4¢/image - // ──────────────────────────────────────────────────────────────────────────── - mod cost { - use super::*; - - mod text { - use super::*; - - fn cost_for_batch(batch: Option) -> u64 { - let request = make_request(Some("p"), None, batch, None); - distill_text(&request).cost.cost_in_usd_cents.unwrap() - } - - #[test] - fn default_batch_costs_4_cents() { assert_eq!(cost_for_batch(None), 4); } - - #[test] - fn batch_of_one_costs_4() { assert_eq!(cost_for_batch(Some(1)), 4); } - - #[test] - fn batch_of_two_costs_8() { assert_eq!(cost_for_batch(Some(2)), 8); } - - #[test] - fn batch_of_three_costs_12() { assert_eq!(cost_for_batch(Some(3)), 12); } - - #[test] - fn batch_of_four_costs_16() { assert_eq!(cost_for_batch(Some(4)), 16); } - - #[test] - fn batch_above_max_clamps_to_16() { assert_eq!(cost_for_batch(Some(7)), 16); } - - #[test] - fn cost_is_independent_of_aspect_ratio() { - let ars = [ - None, - Some(CommonAspectRatio::Square), - Some(CommonAspectRatio::SquareHd), - Some(CommonAspectRatio::WideSixteenByNine), - Some(CommonAspectRatio::WideFourByThree), - Some(CommonAspectRatio::TallNineBySixteen), - Some(CommonAspectRatio::TallThreeByFour), - Some(CommonAspectRatio::Auto), - Some(CommonAspectRatio::Auto2k), - Some(CommonAspectRatio::Auto3k), - Some(CommonAspectRatio::Auto4k), - ]; - for ar in ars { - let request = make_request(Some("p"), ar, Some(2), None); - assert_eq!( - distill_text(&request).cost.cost_in_usd_cents.unwrap(), 8, - "expected 8¢ regardless of aspect ratio {:?}", ar, - ); - } - } - - #[test] - fn cost_is_independent_of_prompt() { - let a = distill_text(&make_request(Some("a cat"), None, Some(1), None)); - let b = distill_text(&make_request(None, None, Some(1), None)); - assert_eq!(a.cost.cost_in_usd_cents, b.cost.cost_in_usd_cents); - } - - #[test] - fn cost_metadata_flags_are_default() { - let d = distill_text(&make_request(Some("p"), None, Some(1), None)); - assert!(!d.cost.is_free); - assert!(!d.cost.is_unlimited); - assert!(!d.cost.is_rate_limited); - assert!(!d.cost.has_watermark); - } - } - - mod edit { - use super::*; - - fn cost_edit(batch: Option, num_refs: usize) -> u64 { - let (tokens, hydration) = fake_image_refs(num_refs); - let request = make_request(Some("p"), None, batch, Some(tokens)); - distill_edit(&request, &hydration).cost.cost_in_usd_cents.unwrap() - } - - #[test] - fn default_batch_costs_4() { assert_eq!(cost_edit(None, 1), 4); } - - #[test] - fn batch_of_one_costs_4() { assert_eq!(cost_edit(Some(1), 1), 4); } - - #[test] - fn batch_of_two_costs_8() { assert_eq!(cost_edit(Some(2), 1), 8); } - - #[test] - fn batch_of_four_costs_16() { assert_eq!(cost_edit(Some(4), 1), 16); } - - #[test] - fn batch_above_max_clamps_to_16() { assert_eq!(cost_edit(Some(9), 1), 16); } - - #[test] - fn cost_independent_of_input_image_count() { - for num_refs in [1usize, 2, 3, 5] { - assert_eq!(cost_edit(Some(2), num_refs), 8, "expected 8¢ regardless of {} input images", num_refs); - } - } - - #[test] - fn edit_and_text_cost_match() { - for batch in [1u16, 2, 3, 4] { - let text = distill_text(&make_request(Some("p"), None, Some(batch), None)) - .cost.cost_in_usd_cents.unwrap(); - let (tokens, hydration) = fake_image_refs(1); - let edit = distill_edit( - &make_request(Some("p"), None, Some(batch), Some(tokens)), - &hydration, - ).cost.cost_in_usd_cents.unwrap(); - assert_eq!(text, edit, "text/edit cost diverged at batch={}", batch); - } - } - - #[test] - fn cost_metadata_flags_are_default() { - let (tokens, hydration) = fake_image_refs(1); - let d = distill_edit(&make_request(Some("p"), None, Some(1), Some(tokens)), &hydration); - assert!(!d.cost.is_free); - assert!(!d.cost.is_unlimited); - assert!(!d.cost.is_rate_limited); - assert!(!d.cost.has_watermark); - } - } - } - - // ──────────────────────────────────────────────────────────────────────────── - // PLAN - // ──────────────────────────────────────────────────────────────────────────── - mod plan { - use super::*; - - use artcraft_router::generate::generate_image::image_generation_plan::ImageGenerationPlan; - use artcraft_router::generate::generate_image::plan::fal::plan_generate_image_fal_seedream_5_lite::{ - FalSeedream5LiteImageSize, FalSeedream5LiteNumImages, PlanFalSeedream5Lite, - }; - - fn with_text_plan( - request: &OmniGenImageCostAndGenerateRequest, - assertion: F, - ) { - let distilled = distill_text(request); - match distilled.plan() { - ImageGenerationPlan::FalSeedream5Lite(plan) => assertion(plan), - other => panic!("expected FalSeedream5Lite, got {:?}", other), - } - } - - fn with_edit_plan( - request: &OmniGenImageCostAndGenerateRequest, - hydration: &HashMap, - assertion: F, - ) { - let distilled = distill_edit(request, hydration); - match distilled.plan() { - ImageGenerationPlan::FalSeedream5Lite(plan) => assertion(plan), - other => panic!("expected FalSeedream5Lite, got {:?}", other), - } - } - - mod text { - use super::*; - - // ── Mode detection ──────────────────────────────────────────────────── - - #[test] - fn text_mode_has_empty_image_urls() { - with_text_plan(&make_request(Some("p"), None, Some(1), None), |plan| { - assert!(plan.image_urls.is_empty()); - }); - } - - // ── Image size: direct mappings ─────────────────────────────────────── - - #[test] - fn default_image_size_is_none() { - with_text_plan(&make_request(Some("p"), None, Some(1), None), |plan| { - assert!(plan.image_size.is_none()); - }); - } - - #[test] - fn square_yields_square() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Square), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::Square))); - }); - } - - #[test] - fn square_hd_yields_square_hd() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::SquareHd), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::SquareHd))); - }); - } - - #[test] - fn wide_16x9_yields_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideSixteenByNine), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_alias_yields_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Wide), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_4x3_yields_landscape_four_three() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideFourByThree), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::LandscapeFourThree))); - }); - } - - #[test] - fn tall_9x16_yields_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallNineBySixteen), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_alias_yields_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Tall), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_3x4_yields_portrait_four_three() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallThreeByFour), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::PortraitFourThree))); - }); - } - - // ── Image size: Auto variants (no bare Auto, no Auto4k) ─────────────── - - #[test] - fn auto_yields_auto_2k() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Auto), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::Auto2k))); - }); - } - - #[test] - fn auto_2k_yields_auto_2k() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Auto2k), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::Auto2k))); - }); - } - - #[test] - fn auto_4k_falls_back_to_auto_3k() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::Auto4k), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::Auto3k))); - }); - } - - // ── Image size: fallbacks ───────────────────────────────────────────── - - #[test] - fn wide_5x4_falls_back_to_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideFiveByFour), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_3x2_falls_back_to_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideThreeByTwo), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn wide_21x9_falls_back_to_landscape_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::WideTwentyOneByNine), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::LandscapeSixteenNine))); - }); - } - - #[test] - fn tall_4x5_falls_back_to_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallFourByFive), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_2x3_falls_back_to_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallTwoByThree), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::PortraitSixteenNine))); - }); - } - - #[test] - fn tall_9x21_falls_back_to_portrait_sixteen_nine() { - with_text_plan(&make_request(Some("p"), Some(CommonAspectRatio::TallNineByTwentyOne), Some(1), None), |plan| { - assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::PortraitSixteenNine))); - }); - } - - // ── Num images ──────────────────────────────────────────────────────── - - #[test] - fn default_batch_count_is_one() { - with_text_plan(&make_request(Some("p"), None, None, None), |plan| { - assert!(matches!(plan.num_images, FalSeedream5LiteNumImages::One)); - }); - } - - #[test] - fn batch_direct_mapping() { - for (count, expected) in [ - (1u16, FalSeedream5LiteNumImages::One), (2, FalSeedream5LiteNumImages::Two), - (3, FalSeedream5LiteNumImages::Three), (4, FalSeedream5LiteNumImages::Four), - ] { - with_text_plan(&make_request(Some("p"), None, Some(count), None), |plan| { - assert!(std::mem::discriminant(&plan.num_images) == std::mem::discriminant(&expected)); - }); - } - } - - #[test] - fn batch_above_four_clamps_to_four() { - with_text_plan(&make_request(Some("p"), None, Some(9), None), |plan| { - assert!(matches!(plan.num_images, FalSeedream5LiteNumImages::Four)); - }); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn prompt_is_passed_through() { - with_text_plan(&make_request(Some("a corgi in a hat"), None, Some(1), None), |plan| { - assert_eq!(plan.prompt, Some("a corgi in a hat".to_string())); - }); - } - - #[test] - fn missing_prompt_is_none() { - with_text_plan(&make_request(None, None, Some(1), None), |plan| { - assert_eq!(plan.prompt, None); - }); - } - } - - mod edit { - use super::*; - - // ── Mode detection / image-url passthrough ──────────────────────────── - - #[test] - fn edit_mode_populates_image_urls() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { - assert_eq!(plan.image_urls.len(), 1); - assert!(plan.image_urls[0].starts_with("https://fake.example.com/")); - }, - ); - } - - #[test] - fn edit_mode_with_three_image_refs() { - let (tokens, hydration) = fake_image_refs(3); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.image_urls.len(), 3); }, - ); - } - - #[test] - fn edit_mode_with_five_image_refs() { - let (tokens, hydration) = fake_image_refs(5); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.image_urls.len(), 5); }, - ); - } - - #[test] - fn edit_image_urls_are_hydrated_from_map() { - let (tokens, hydration) = fake_image_refs(3); - let expected: Vec = tokens.iter() - .map(|t| hydration.get(t).unwrap().to_string()) - .collect(); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens.clone())), - &hydration, - |plan| { assert_eq!(plan.image_urls, expected); }, - ); - } - - // ── Auto in edit mode also maps to Auto2k (no bare Auto in v5 lite) ─── - - #[test] - fn auto_in_edit_mode_yields_auto_2k() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Auto), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::Auto2k))); }, - ); - } - - #[test] - fn auto_4k_in_edit_mode_yields_auto_3k() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Auto4k), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::Auto3k))); }, - ); - } - - // ── Image size direct mappings in edit mode ─────────────────────────── - - #[test] - fn edit_default_image_size_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert!(plan.image_size.is_none()); }, - ); - } - - #[test] - fn edit_square_yields_square() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Square), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::Square))); }, - ); - } - - #[test] - fn edit_wide_16x9_yields_landscape_sixteen_nine() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::WideSixteenByNine), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::LandscapeSixteenNine))); }, - ); - } - - #[test] - fn edit_tall_9x16_yields_portrait_sixteen_nine() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::TallNineBySixteen), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::PortraitSixteenNine))); }, - ); - } - - #[test] - fn edit_auto_2k() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), Some(CommonAspectRatio::Auto2k), Some(1), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.image_size, Some(FalSeedream5LiteImageSize::Auto2k))); }, - ); - } - - // ── Num images in edit mode ─────────────────────────────────────────── - - #[test] - fn edit_default_batch_count_is_one() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, None, Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalSeedream5LiteNumImages::One)); }, - ); - } - - #[test] - fn edit_batch_of_four_yields_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(4), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalSeedream5LiteNumImages::Four)); }, - ); - } - - #[test] - fn edit_batch_above_four_clamps_to_four() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("p"), None, Some(9), Some(tokens)), - &hydration, - |plan| { assert!(matches!(plan.num_images, FalSeedream5LiteNumImages::Four)); }, - ); - } - - // ── Prompt passthrough ──────────────────────────────────────────────── - - #[test] - fn edit_prompt_is_passed_through() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(Some("make it shiny"), None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, Some("make it shiny".to_string())); }, - ); - } - - #[test] - fn edit_missing_prompt_is_none() { - let (tokens, hydration) = fake_image_refs(1); - with_edit_plan( - &make_request(None, None, Some(1), Some(tokens)), - &hydration, - |plan| { assert_eq!(plan.prompt, None); }, - ); - } - } - } -} diff --git a/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_map.rs b/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_map.rs index 1a52a3b64e..48306c5219 100644 --- a/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_map.rs +++ b/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_map.rs @@ -12,6 +12,7 @@ use std::collections::{HashMap, HashSet}; use std::iter::FromIterator; use tokens::tokens::media_files::MediaFileToken; +#[deprecated(note="Try to use `lookup_media_files_as_cdn_url_list_and_map()` instead")] pub async fn lookup_image_urls_as_map( http_request: &HttpRequest, mysql_connection: &mut PoolConnection, diff --git a/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_optional_list.rs b/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_optional_list.rs index 243a782309..99bade93a7 100644 --- a/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_optional_list.rs +++ b/crates/service/web/storyteller_web/src/util/lookup/lookup_image_urls_as_optional_list.rs @@ -12,6 +12,7 @@ use std::collections::HashSet; use std::iter::FromIterator; use tokens::tokens::media_files::MediaFileToken; +#[deprecated(note="I don't think we should be returning Option based on an empty() check at the end. This is weird.")] pub async fn lookup_image_urls_as_optional_list( http_request: &HttpRequest, mysql_connection: &mut PoolConnection, diff --git a/crates/service/web/storyteller_web/src/util/lookup/lookup_media_file_urls_as_map.rs b/crates/service/web/storyteller_web/src/util/lookup/lookup_media_file_urls_as_map.rs index d9bcdc98ec..8e516f243d 100644 --- a/crates/service/web/storyteller_web/src/util/lookup/lookup_media_file_urls_as_map.rs +++ b/crates/service/web/storyteller_web/src/util/lookup/lookup_media_file_urls_as_map.rs @@ -13,6 +13,7 @@ use std::iter::FromIterator; use tokens::tokens::media_files::MediaFileToken; use url::Url; +#[deprecated(note="Try to use `lookup_media_files_as_cdn_url_list_and_map()` instead")] pub async fn lookup_media_file_urls_as_map( http_request: &HttpRequest, mysql_connection: &mut PoolConnection, diff --git a/crates/service/web/storyteller_web/src/util/lookup/lookup_media_files_as_cdn_url_list_and_map.rs b/crates/service/web/storyteller_web/src/util/lookup/lookup_media_files_as_cdn_url_list_and_map.rs new file mode 100644 index 0000000000..af67066bf1 --- /dev/null +++ b/crates/service/web/storyteller_web/src/util/lookup/lookup_media_files_as_cdn_url_list_and_map.rs @@ -0,0 +1,96 @@ +use std::collections::{HashMap, HashSet}; +use std::iter::FromIterator; + +use actix_web::HttpRequest; +use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath; +use log::{error, warn}; +use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::batch_get_media_files_by_tokens_with_connection; +use server_environment::ServerEnvironment; +use sqlx::pool::PoolConnection; +use sqlx::MySql; +use tokens::tokens::media_files::MediaFileToken; + +use crate::http_server::common_responses::common_web_error::CommonWebError; +use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder; +use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain; + +pub struct MediaFilesAsCdnUrlListAndMap { + /// This is an in-order list of all the media files as their asset URLs. + /// This is in the order they were requested. + pub ordered_url_list: Vec, + + /// This is an unordered map of media file token to asset URL. + pub token_to_url_map: HashMap, +} + +pub async fn lookup_media_files_as_cdn_url_list_and_map( + http_request: &HttpRequest, + mysql_connection: &mut PoolConnection, + server_environment: ServerEnvironment, + tokens: &[MediaFileToken], +) -> Result { + const CAN_SEE_DELETED: bool = false; + + if tokens.is_empty() { + return Ok(MediaFilesAsCdnUrlListAndMap { + ordered_url_list: Vec::new(), + token_to_url_map: HashMap::new(), + }); + } + + let result = batch_get_media_files_by_tokens_with_connection( + mysql_connection, + tokens, + CAN_SEE_DELETED, + ).await; + + let media_files = match result { + Ok(files) => files, + Err(err) => { + error!("Error getting media files by tokens: {:?}", err); + return Err(CommonWebError::ServerError); + } + }; + + if media_files.len() != tokens.len() { + warn!("Wrong number of media files returned for tokens: {} found for {} tokens", media_files.len(), tokens.len()); + + let requested: HashSet<&MediaFileToken> = HashSet::from_iter(tokens.iter()); + let returned: HashSet<&MediaFileToken> = HashSet::from_iter(media_files.iter().map(|m| &m.token)); + + let diff = requested.difference(&returned) + .cloned() + .collect::>(); + + return Err(CommonWebError::BadInputWithSimpleMessage( + format!("Not all media files could be found. Media files found: {}, tokens provided: {}, in original: {:?}, req {:?}, ret {:?}", + media_files.len(), tokens.len(), diff, requested, returned))); + } + + let media_domain = get_media_domain(http_request); + + let mut ordered_url_list = Vec::with_capacity(media_files.len()); + let mut token_to_url_map = HashMap::with_capacity(media_files.len()); + + for file in media_files { + let public_bucket_path = MediaFileBucketPath::from_object_hash( + &file.public_bucket_directory_hash, + file.maybe_public_bucket_prefix.as_deref(), + file.maybe_public_bucket_extension.as_deref()); + + let media_links = MediaLinksBuilder::from_media_path_and_env( + media_domain, + server_environment, + &public_bucket_path); + + let url = media_links.cdn_url.to_string(); + + ordered_url_list.push(url.clone()); + token_to_url_map.insert(file.token, url); + } + + Ok(MediaFilesAsCdnUrlListAndMap { + ordered_url_list, + token_to_url_map, + }) +} diff --git a/crates/service/web/storyteller_web/src/util/lookup/mod.rs b/crates/service/web/storyteller_web/src/util/lookup/mod.rs index d1e482e773..41b796b4f3 100644 --- a/crates/service/web/storyteller_web/src/util/lookup/mod.rs +++ b/crates/service/web/storyteller_web/src/util/lookup/mod.rs @@ -1,5 +1,6 @@ pub mod fetch_all_required_media_files; pub mod lookup_image_urls_as_optional_list; pub mod lookup_image_urls_as_map; +pub mod lookup_media_files_as_cdn_url_list_and_map; pub mod lookup_media_file_urls_as_map; -pub mod resolve_referral_info; \ No newline at end of file +pub mod resolve_referral_info; diff --git a/frontend/apps/artcraft-webapp/src/main.tsx b/frontend/apps/artcraft-webapp/src/main.tsx index df28a4f27f..72e9caa1e0 100644 --- a/frontend/apps/artcraft-webapp/src/main.tsx +++ b/frontend/apps/artcraft-webapp/src/main.tsx @@ -19,7 +19,7 @@ if (import.meta.env.DEV) { window.location.origin, ); // NB: This is for Brandon to test with storyteller-web locally: - //StorytellerApiHostStore.getInstance().setDevelopment(); + StorytellerApiHostStore.getInstance().setDevelopment(); } catch (e) { console.warn("Failed to set dev API host override", e); } diff --git a/frontend/apps/artcraft-website/src/main.tsx b/frontend/apps/artcraft-website/src/main.tsx index df28a4f27f..72e9caa1e0 100644 --- a/frontend/apps/artcraft-website/src/main.tsx +++ b/frontend/apps/artcraft-website/src/main.tsx @@ -19,7 +19,7 @@ if (import.meta.env.DEV) { window.location.origin, ); // NB: This is for Brandon to test with storyteller-web locally: - //StorytellerApiHostStore.getInstance().setDevelopment(); + StorytellerApiHostStore.getInstance().setDevelopment(); } catch (e) { console.warn("Failed to set dev API host override", e); }