//! End-to-end checks for Enhance (#192) against a mock ComfyUI: nothing here //! touches a real ComfyUI, a real library or the prod DB. use std::collections::HashMap; use std::path::PathBuf; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use actix_web::test as awtest; use actix_web::web::Data; use actix_web::{App, HttpResponse, HttpServer, web}; use diesel::prelude::*; use diesel::sql_query; use diesel::sql_types::Text; use diesel::sqlite::SqliteConnection; use serde_json::{Value, json}; use super::config::{EnhanceConfig, PathMapping, PromptMode, TemplateConfig}; use super::service::EnhanceService; use super::template::Template; use crate::data::Claims; use crate::database::enhance_dao::{EnhanceJobDao, SqliteEnhanceJobDao}; use crate::database::test::in_memory_db_connection; use crate::database::{ExifDao, SqliteExifDao, SqliteVersionDao}; use crate::state::AppState; use crate::versions::xmp::test_png::{producer_xmp, write_png}; const RESTORE: &str = include_str!("../../tests/fixtures/enhance/restore_turbo.json"); const EDIT: &str = include_str!("../../tests/fixtures/enhance/edit_turbo.json"); // ---- mock ComfyUI --------------------------------------------------------- #[derive(Default)] struct Mock { submitted: Mutex>, queue_running: Mutex>, queue_pending: Mutex>, /// prompt_id -> history entry history: Mutex>, newest: Mutex>, approve_calls: Mutex>, approve_path: Mutex, approve_gone: Mutex, free_calls: AtomicUsize, reject_prompts: Mutex>, view_404: Mutex, } async fn start_mock() -> (Arc, String) { let mock = Arc::new(Mock::default()); let m = mock.clone(); let server = HttpServer::new(move || { App::new() .app_data(Data::from(m.clone())) .route("/system_stats", web::get().to(|| async { HttpResponse::Ok().json(json!({})) })) .route( "/prompt", web::post().to(|m: Data, body: web::Json| async move { if let Some(err) = m.reject_prompts.lock().unwrap().clone() { return HttpResponse::BadRequest().json(err); } let mut s = m.submitted.lock().unwrap(); s.push(body.into_inner()); let id = format!("prompt-{}", s.len()); m.queue_pending.lock().unwrap().push(id.clone()); HttpResponse::Ok().json(json!({"prompt_id": id, "number": s.len(), "node_errors": {}})) }), ) .route( "/queue", web::get().to(|m: Data| async move { let item = |id: &String| json!([1, id, {}, {}, []]); HttpResponse::Ok().json(json!({ "queue_running": m.queue_running.lock().unwrap().iter().map(item).collect::>(), "queue_pending": m.queue_pending.lock().unwrap().iter().map(item).collect::>(), })) }), ) .route( "/history", web::get().to(|m: Data| async move { match m.newest.lock().unwrap().clone() { Some(e) => HttpResponse::Ok().json(json!({"newest": e})), None => HttpResponse::Ok().json(json!({})), } }), ) .route( "/history/{id}", web::get().to(|m: Data, id: web::Path| async move { match m.history.lock().unwrap().get(id.as_str()) { Some(e) => HttpResponse::Ok().json(json!({ id.as_str(): e })), None => HttpResponse::Ok().json(json!({})), } }), ) .route( "/view", web::get().to(|m: Data, q: web::Query>| async move { if *m.view_404.lock().unwrap() { return HttpResponse::NotFound().finish(); } assert_eq!(q.get("type").map(String::as_str), Some("temp")); HttpResponse::Ok() .content_type("image/png") .body(format!("PNG:{}", q["filename"])) }), ) .route( "/imageapi/approve", web::post().to(|m: Data, body: web::Json| async move { if *m.approve_gone.lock().unwrap() { return HttpResponse::NotFound().json(json!({"error": "no longer pending"})); } m.approve_calls .lock() .unwrap() .push(body["id"].as_str().unwrap().to_string()); HttpResponse::Ok().json(json!({"path": m.approve_path.lock().unwrap().clone()})) }), ) .route( "/free", web::post().to(|m: Data, body: web::Json| async move { assert_eq!(body["unload_models"], true); m.free_calls.fetch_add(1, Ordering::SeqCst); HttpResponse::Ok().finish() }), ) }) .workers(1) .bind(("127.0.0.1", 0)) .unwrap(); let addr = server.addrs()[0]; actix_rt::spawn(server.run()); (mock, format!("http://{addr}")) } fn finished_entry(pending_id: &str, ts: i64) -> Value { json!({ "status": {"status_str": "success", "completed": true, "messages": [["execution_start", {"timestamp": ts - 3000}], ["execution_success", {"timestamp": ts}]]}, "outputs": {"36": { "images": [{"filename": format!("imageapi_pending_{pending_id}.png"), "subfolder": "", "type": "temp"}], "imageapi_versions": [{"id": pending_id, "saved_path": null}], }}, }) } // ---- fixture -------------------------------------------------------------- struct Fx { mock: Arc, svc: Data, state: Data, exif: Data>>, conn: Arc>, root: PathBuf, } fn template(name: &str, mode: PromptMode, json: &str) -> Template { Template::from_value( TemplateConfig { name: name.into(), label: name.to_uppercase(), description: format!("the {name} template"), file: "x.json".into(), prompt_mode: mode, prompt_join: " ".into(), output_megapixels: 1.0, }, serde_json::from_str(json).unwrap(), ) .unwrap() } async fn fx_with(url: Option, tweak: impl FnOnce(&mut EnhanceConfig)) -> Fx { let (mock, mock_url) = start_mock().await; let mut state = AppState::test_state(); let conn = Arc::new(Mutex::new(in_memory_db_connection())); state.version_dao = Arc::new(Mutex::new(Box::new(SqliteVersionDao::from_shared( conn.clone(), )))); let library = state.primary_library().clone(); let root = PathBuf::from(&library.root_path); sql_query("UPDATE libraries SET root_path = ? WHERE id = 1") .bind::(&library.root_path) .execute(&mut *conn.lock().unwrap()) .unwrap(); let mut cfg = EnhanceConfig { url: Some(url.unwrap_or(mock_url)), path_map: vec![PathMapping { from: format!("{}/", library.root_path), to: "Z:\\".into(), }], ..EnhanceConfig::default() }; tweak(&mut cfg); let dao: Arc>> = Arc::new(Mutex::new(Box::new( SqliteEnhanceJobDao::from_shared(conn.clone()), ))); let svc = EnhanceService::new( cfg, vec![ template("restore", PromptMode::Append, RESTORE), template("edit", PromptMode::Replace, EDIT), ], dao, ) .unwrap(); Fx { mock, svc: Data::new(svc), exif: Data::new(Mutex::new( Box::new(SqliteExifDao::from_shared(conn.clone())) as Box, )), state: Data::new(state), conn, root, } } async fn fx() -> Fx { fx_with(None, |_| {}).await } impl Fx { fn put(&self, rel: &str, bytes: &[u8]) -> String { let abs = self.root.join(rel); std::fs::create_dir_all(abs.parent().unwrap()).unwrap(); std::fs::write(&abs, bytes).unwrap(); self.register(rel) } fn register(&self, rel: &str) -> String { let hash = crate::content_hash::compute(&self.root.join(rel)) .unwrap() .content_hash; sql_query( "INSERT INTO image_exif (library_id, rel_path, created_time, last_modified, content_hash) \ VALUES (1, ?, 0, 0, ?)", ) .bind::(rel) .bind::(&hash) .execute(&mut *self.conn.lock().unwrap()) .unwrap(); hash } fn token() -> String { jsonwebtoken::encode( &jsonwebtoken::Header::default(), &Claims::valid_user("1".to_string()), &jsonwebtoken::EncodingKey::from_secret(b"test_key"), ) .unwrap() } async fn call(&self, method: &str, uri: &str, body: Option) -> (u16, Value, Vec) { let app = awtest::init_service( App::new() .service(super::get_templates) .service(super::submit) .service(super::list_jobs) .service(super::get_job) .service(super::preview) .service(super::approve) .service(super::discard) .service(crate::versions::get_versions) .app_data(self.svc.clone()) .app_data(self.state.clone()) .app_data(self.exif.clone()), ) .await; let mut req = match method { "POST" => awtest::TestRequest::post(), _ => awtest::TestRequest::get(), } .uri(uri) .insert_header(("Authorization", format!("Bearer {}", Fx::token()))); if let Some(b) = body { req = req.set_json(b); } let resp = awtest::call_service(&app, req.to_request()).await; let status = resp.status().as_u16(); let bytes = awtest::read_body(resp).await.to_vec(); let json = serde_json::from_slice(&bytes).unwrap_or(Value::Null); (status, json, bytes) } } async fn call(fx: &Fx, method: &str, uri: &str, body: Option) -> (u16, Value, Vec) { fx.call(method, uri, body).await } const PHOTO: &str = "2009_200/019_19.JPG"; // ---- tests ---------------------------------------------------------------- #[actix_rt::test] async fn templates_report_availability_and_shape() { let fx = fx().await; let (st, body, _) = call(&fx, "GET", "/enhance/templates", None).await; assert_eq!(st, 200); assert_eq!(body["available"], true); assert!(body.get("reason").is_none()); assert_eq!(body["default"], "restore"); let t = &body["templates"]; assert_eq!(t[0]["name"], "restore"); assert_eq!(t[0]["label"], "RESTORE"); assert_eq!(t[0]["accepts_prompt"], true); assert_eq!(t[0]["prompt_mode"], "append"); assert_eq!(t[0]["prompt_required"], false); assert_eq!(t[0]["accepts_reference"], false); assert_eq!(t[0]["default"], true); assert_eq!(t[0]["output_megapixels"], 1.0); assert_eq!(t[0]["warn_above_megapixels"], 1.5); assert_eq!(t[1]["name"], "edit"); assert_eq!(t[1]["prompt_mode"], "replace"); assert_eq!(t[1]["prompt_required"], true); assert_eq!(t[1]["accepts_reference"], true); assert_eq!(t[1]["default"], false); } #[actix_rt::test] async fn comfy_down_means_unavailable_and_a_clear_submit_error() { // Nothing listens on port 9 (discard) of loopback. let fx = fx_with(Some("http://127.0.0.1:9".into()), |_| {}).await; fx.put(PHOTO, b"photo"); let (st, body, _) = call(&fx, "GET", "/enhance/templates", None).await; assert_eq!(st, 200); assert_eq!(body["available"], false); assert_eq!(body["reason"], "ComfyUI isn't running"); assert_eq!(body["templates"].as_array().unwrap().len(), 2); let (st, body, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; assert_eq!(st, 503); assert_eq!(body["error"], "ComfyUI isn't running"); assert_eq!(body["code"], "unavailable"); } #[actix_rt::test] async fn unconfigured_enhance_is_off_not_an_error() { let fx = fx_with(None, |c| c.url = None).await; let (st, body, _) = call(&fx, "GET", "/enhance/templates", None).await; assert_eq!((st, &body["available"]), (200, &json!(false))); let (st, _, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; assert_eq!(st, 503); } #[actix_rt::test] async fn submit_maps_the_path_and_uses_a_fresh_seed_each_time() { let fx = fx().await; fx.put(PHOTO, b"photo"); let body = json!({"path": PHOTO, "template": "restore", "prompt": " warmer colours "}); let (st, job1, _) = call(&fx, "POST", "/enhance", Some(body.clone())).await; assert_eq!(st, 200, "{job1}"); let (_, job2, _) = call(&fx, "POST", "/enhance", Some(body)).await; assert_eq!(job1["status"], "queued"); assert_eq!(job1["path"], PHOTO); assert_eq!(job1["template"], "restore"); assert_eq!(job1["prompt"], "warmer colours"); assert_eq!(job1["pending"], json!([])); assert!(job1["created_at"].as_i64().unwrap() > 0); assert_ne!(job1["id"], job2["id"]); let sent = fx.mock.submitted.lock().unwrap(); assert_eq!(sent.len(), 2); let p1 = &sent[0]["prompt"]; let p2 = &sent[1]["prompt"]; assert_eq!(p1["37"]["inputs"]["path"], "Z:\\2009_200\\019_19.JPG"); assert_eq!(p1["36"]["inputs"]["require_approval"], true); assert!( p1["13"]["inputs"]["value"] .as_str() .unwrap() .ends_with("overlays. warmer colours") ); assert_ne!(p1["5"]["inputs"]["seed"], p2["5"]["inputs"]["seed"]); assert_ne!( p1["14"]["inputs"]["sampling_mode.seed"], p2["14"]["inputs"]["sampling_mode.seed"] ); assert_eq!(sent[0]["client_id"], "imageapi"); } #[actix_rt::test] async fn submit_validation() { let fx = fx().await; fx.put(PHOTO, b"photo"); // Edit needs an instruction. let (st, body, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "edit"})), ) .await; assert_eq!(st, 400, "{body}"); assert!(body["error"].as_str().unwrap().contains("instruction")); let (st, _, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "edit", "prompt": " "})), ) .await; assert_eq!(st, 400); // Unknown template / photo / traversal / reference on a template without one. let (st, _, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "nope"})), ) .await; assert_eq!(st, 404); let (st, _, _) = call( &fx, "POST", "/enhance", Some(json!({"path": "missing.jpg", "template": "restore"})), ) .await; assert_eq!(st, 404); let (st, _, _) = call( &fx, "POST", "/enhance", Some(json!({"path": "../etc/passwd", "template": "restore"})), ) .await; assert_eq!(st, 400); fx.put("ref.jpg", b"other"); let (st, _, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore", "reference_path": "ref.jpg"})), ) .await; assert_eq!(st, 400); assert!( fx.mock.submitted.lock().unwrap().is_empty(), "nothing reached ComfyUI" ); } #[actix_rt::test] async fn edit_without_reference_drops_the_node_and_with_one_keeps_it() { let fx = fx().await; fx.put(PHOTO, b"photo"); fx.put("ref/other.jpg", b"other view"); let (st, _, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "edit", "prompt": "make it b&w"})), ) .await; assert_eq!(st, 200); let (st, job, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "edit", "prompt": "x", "reference_path": "ref/other.jpg"})), ) .await; assert_eq!(st, 200); assert_eq!(job["reference_path"], "ref/other.jpg"); let sent = fx.mock.submitted.lock().unwrap(); let without = &sent[0]["prompt"]; assert!(without.get("38").is_none()); assert!(!without.to_string().contains("[\"38\"")); assert_eq!(without["13"]["inputs"]["value"], "make it b&w"); let with = &sent[1]["prompt"]; assert_eq!(with["38"]["inputs"]["path"], "Z:\\ref\\other.jpg"); } #[actix_rt::test] async fn comfy_rejection_is_reported_and_a_missing_reference_is_named() { let fx = fx().await; fx.put(PHOTO, b"photo"); *fx.mock.reject_prompts.lock().unwrap() = Some(json!({ "error": {"type": "prompt_outputs_failed_validation", "message": "Prompt outputs failed validation", "details": ""}, "node_errors": {"36": {"class_type": "ImageApiSaveImageVersion", "errors": [{"message": "Required input is missing"}]}}, })); let (st, body, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "edit", "prompt": "x"})), ) .await; assert_eq!(st, 422); assert!( body["error"] .as_str() .unwrap() .contains("needs a reference photo"), "{body}" ); // Restore never dropped a reference, so the message is ComfyUI's own. let (st, body, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; assert_eq!(st, 422); assert!(!body["error"].as_str().unwrap().contains("reference")); assert!( body["error"] .as_str() .unwrap() .contains("Required input is missing") ); } #[actix_rt::test] async fn job_moves_queued_running_ready_and_the_preview_is_relayed() { let fx = fx().await; fx.put(PHOTO, b"photo"); let (_, job, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; let id = job["id"].as_str().unwrap().to_string(); // Behind someone else's run: queued with a position. fx.mock.queue_running.lock().unwrap().push("other".into()); fx.svc.poll_once().await; let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; assert_eq!( (j["status"].as_str().unwrap(), j["queue_position"].as_i64()), ("queued", Some(1)) ); // Now it runs. *fx.mock.queue_running.lock().unwrap() = vec!["prompt-1".into()]; fx.mock.queue_pending.lock().unwrap().clear(); fx.svc.poll_once().await; let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; assert_eq!(j["status"], "running"); // And finishes. fx.mock.queue_running.lock().unwrap().clear(); fx.mock .history .lock() .unwrap() .insert("prompt-1".into(), finished_entry("abc", 10_000)); fx.svc.poll_once().await; let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; assert_eq!(j["status"], "ready"); assert!(j.get("queue_position").is_none()); assert_eq!(j["pending"].as_array().unwrap().len(), 1); assert_eq!(j["pending"][0]["n"], 0); assert_eq!(j["pending"][0]["approved"], false); // The staged file name is ComfyUI-internal and not exposed. assert!(!j.to_string().contains("imageapi_pending")); let uri = j["pending"][0]["preview"].as_str().unwrap().to_string(); let (st, _, bytes) = call(&fx, "GET", &uri, None).await; assert_eq!(st, 200); assert_eq!(bytes, b"PNG:imageapi_pending_abc.png"); let (st, _, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}/preview/5"), None).await; assert_eq!(st, 404); // The photo's job list shows it (and a failed one too). let (_, list, _) = call(&fx, "GET", &format!("/enhance/jobs?path={PHOTO}"), None).await; assert_eq!(list["jobs"].as_array().unwrap().len(), 1); assert_eq!(list["jobs"][0]["id"], id.as_str()); } #[actix_rt::test] async fn failed_run_reports_comfys_message() { let fx = fx().await; fx.put(PHOTO, b"photo"); let (_, job, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; fx.mock.queue_pending.lock().unwrap().clear(); fx.mock.history.lock().unwrap().insert( "prompt-1".into(), json!({"status": {"status_str": "error", "completed": false, "messages": [ ["execution_error", {"exception_message": "CUDA out of memory\n", "timestamp": 5}]]}, "outputs": {}}), ); fx.svc.poll_once().await; let (_, j, _) = call( &fx, "GET", &format!("/enhance/jobs/{}", job["id"].as_str().unwrap()), None, ) .await; assert_eq!(j["status"], "failed"); assert_eq!(j["error"], "CUDA out of memory"); } #[actix_rt::test] async fn a_job_comfy_forgot_expires() { let fx = fx().await; fx.put(PHOTO, b"photo"); let (_, job, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; // ComfyUI restarted: empty queue, empty history. fx.mock.queue_pending.lock().unwrap().clear(); fx.svc.poll_once().await; let (_, j, _) = call( &fx, "GET", &format!("/enhance/jobs/{}", job["id"].as_str().unwrap()), None, ) .await; assert_eq!(j["status"], "expired"); assert!(j["error"].as_str().unwrap().contains("restarted")); assert_eq!(j["pending"], json!([])); } #[actix_rt::test] async fn ready_job_expires_when_the_preview_is_gone_and_approve_says_so() { let fx = fx().await; fx.put(PHOTO, b"photo"); let (_, job, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; let id = job["id"].as_str().unwrap().to_string(); fx.mock.queue_pending.lock().unwrap().clear(); fx.mock .history .lock() .unwrap() .insert("prompt-1".into(), finished_entry("abc", 1)); fx.svc.poll_once().await; *fx.mock.view_404.lock().unwrap() = true; let (st, body, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}/preview/0"), None).await; assert_eq!((st, body["code"].as_str()), (410, Some("expired"))); let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; assert_eq!(j["status"], "expired"); let (st, _, _) = call( &fx, "POST", &format!("/enhance/jobs/{id}/approve"), Some(json!({"n": 0})), ) .await; assert_eq!(st, 410); assert!(fx.mock.approve_calls.lock().unwrap().is_empty()); } #[actix_rt::test] async fn approve_gone_on_comfy_side_expires_the_job() { let fx = fx().await; fx.put(PHOTO, b"photo"); let (_, job, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; let id = job["id"].as_str().unwrap().to_string(); fx.mock.queue_pending.lock().unwrap().clear(); fx.mock .history .lock() .unwrap() .insert("prompt-1".into(), finished_entry("abc", 1)); fx.svc.poll_once().await; *fx.mock.approve_gone.lock().unwrap() = true; let (st, _, _) = call( &fx, "POST", &format!("/enhance/jobs/{id}/approve"), Some(json!({"n": 0})), ) .await; assert_eq!(st, 410); let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; assert_eq!(j["status"], "expired"); } #[actix_rt::test] async fn approve_ingests_links_and_is_idempotent() { let fx = fx().await; let original = fx.put(PHOTO, b"the original photo"); let (_, job, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; let id = job["id"].as_str().unwrap().to_string(); fx.mock.queue_pending.lock().unwrap().clear(); fx.mock .history .lock() .unwrap() .insert("prompt-1".into(), finished_entry("abc", 1)); fx.svc.poll_once().await; // What ComfyUI's approve would have written: a PNG with the XMP link. let saved_rel = "2009_200/aiedits/019_19.JPG.restored.png"; let xmp = producer_xmp(&original, "restored", &[]); write_png(&fx.root.join(saved_rel), Some(&xmp), None, 7); *fx.mock.approve_path.lock().unwrap() = "Z:\\2009_200\\aiedits\\019_19.JPG.restored.png".into(); let (st, v, _) = call( &fx, "POST", &format!("/enhance/jobs/{id}/approve"), Some(json!({"n": 0})), ) .await; assert_eq!(st, 200, "{v}"); assert_eq!(v["path"], saved_rel); assert_eq!(v["library_id"], 1); assert_eq!(v["kind"], "restored"); assert_eq!(v["source"], "xmp"); for k in ["hash", "created_at", "references", "width", "height"] { assert!(v.get(k).is_some(), "missing {k}: {v}"); } assert_eq!(fx.mock.approve_calls.lock().unwrap().as_slice(), ["abc"]); // Visible in /image/versions right away. let (_, vs, _) = call(&fx, "GET", &format!("/image/versions?path={PHOTO}"), None).await; assert_eq!(vs["versions"].as_array().unwrap().len(), 1); assert_eq!(vs["versions"][0]["path"], saved_rel); // Approving again: same version, ComfyUI not asked, no duplicate row. let (st, v2, _) = call( &fx, "POST", &format!("/enhance/jobs/{id}/approve"), Some(json!({"n": 0})), ) .await; assert_eq!(st, 200); assert_eq!(v2["hash"], v["hash"]); assert_eq!(fx.mock.approve_calls.lock().unwrap().len(), 1); let (_, vs, _) = call(&fx, "GET", &format!("/image/versions?path={PHOTO}"), None).await; assert_eq!(vs["versions"].as_array().unwrap().len(), 1); let rows: i64 = { #[derive(QueryableByName)] struct N { #[diesel(sql_type = diesel::sql_types::BigInt)] n: i64, } sql_query("SELECT COUNT(*) AS n FROM image_exif WHERE rel_path = ?") .bind::(saved_rel) .get_result::(&mut *fx.conn.lock().unwrap()) .unwrap() .n }; assert_eq!(rows, 1); let (_, j, _) = call(&fx, "GET", &format!("/enhance/jobs/{id}"), None).await; assert_eq!(j["status"], "approved"); assert_eq!(j["pending"][0]["approved"], true); assert_eq!(j["saved"][0]["path"], saved_rel); } #[actix_rt::test] async fn approve_of_a_path_outside_every_library_is_reported_not_lost() { let fx = fx().await; fx.put(PHOTO, b"photo"); let (_, job, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; let id = job["id"].as_str().unwrap().to_string(); fx.mock.queue_pending.lock().unwrap().clear(); fx.mock .history .lock() .unwrap() .insert("prompt-1".into(), finished_entry("abc", 1)); fx.svc.poll_once().await; *fx.mock.approve_path.lock().unwrap() = "C:\\somewhere\\else.png".into(); let (st, body, _) = call( &fx, "POST", &format!("/enhance/jobs/{id}/approve"), Some(json!({"n": 0})), ) .await; assert_eq!(st, 502); assert!( body["error"] .as_str() .unwrap() .contains("C:\\somewhere\\else.png") ); } #[actix_rt::test] async fn approve_before_ready_is_a_400() { let fx = fx().await; fx.put(PHOTO, b"photo"); let (_, job, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; let (st, _, _) = call( &fx, "POST", &format!("/enhance/jobs/{}/approve", job["id"].as_str().unwrap()), Some(json!({"n": 0})), ) .await; assert_eq!(st, 400); let (st, _, _) = call( &fx, "POST", "/enhance/jobs/nope/approve", Some(json!({"n": 0})), ) .await; assert_eq!(st, 404); } #[actix_rt::test] async fn discard_hides_the_job_from_the_photo_list_only() { let fx = fx().await; fx.put(PHOTO, b"photo"); let (_, a, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; let (_, b, _) = call( &fx, "POST", "/enhance", Some(json!({"path": PHOTO, "template": "restore"})), ) .await; let (st, _, _) = call( &fx, "POST", &format!("/enhance/jobs/{}/discard", a["id"].as_str().unwrap()), None, ) .await; assert_eq!(st, 200); let (_, list, _) = call(&fx, "GET", &format!("/enhance/jobs?path={PHOTO}"), None).await; let ids: Vec<&str> = list["jobs"] .as_array() .unwrap() .iter() .map(|j| j["id"].as_str().unwrap()) .collect(); assert_eq!(ids, [b["id"].as_str().unwrap()]); // Still fetchable by id. let (st, _, _) = call( &fx, "GET", &format!("/enhance/jobs/{}", a["id"].as_str().unwrap()), None, ) .await; assert_eq!(st, 200); let (st, _, _) = call(&fx, "POST", "/enhance/jobs/nope/discard", None).await; assert_eq!(st, 404); } #[actix_rt::test] async fn unreachable_active_jobs_become_unavailable_after_the_grace_period() { // Submit against the mock, then point nothing at it: use a service whose // client targets a dead port but whose DB holds an active job. let fx = fx_with(Some("http://127.0.0.1:9".into()), |c| { c.unreachable_grace_secs = 1 }) .await; let now = chrono::Utc::now().timestamp(); let dao_job = crate::database::enhance_dao::EnhanceJobRow { id: "j1".into(), original_hash: "h".into(), library_id: 1, rel_path: "a.jpg".into(), template: "restore".into(), prompt: None, reference_hash: None, reference_path: None, comfy_prompt_id: "p".into(), status: "running".into(), error: None, queue_position: None, pending: "[]".into(), saved: "[]".into(), discarded: 0, created_at: now, updated_at: now, }; SqliteEnhanceJobDao::from_shared(fx.conn.clone()) .insert(&dao_job) .unwrap(); fx.svc.poll_once().await; assert_eq!( fx.svc.get_job("j1").unwrap().status, "running", "within grace" ); tokio::time::sleep(std::time::Duration::from_millis(1100)).await; fx.svc.poll_once().await; let j = fx.svc.get_job("j1").unwrap(); assert_eq!(j.status, "unavailable"); assert_eq!(j.error.as_deref(), Some("ComfyUI isn't running")); } #[actix_rt::test] async fn idle_comfy_is_freed_once_and_activity_rearms_it() { let fx = fx_with(None, |c| c.idle_free_secs = 1).await; *fx.mock.newest.lock().unwrap() = Some(finished_entry("x", 100)); fx.svc.idle_check().await; // baseline assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 0); tokio::time::sleep(std::time::Duration::from_millis(1100)).await; fx.svc.idle_check().await; fx.svc.idle_check().await; assert_eq!( fx.mock.free_calls.load(Ordering::SeqCst), 1, "freed exactly once" ); // Someone runs a workflow in ComfyUI's own page: newest history changes. *fx.mock.newest.lock().unwrap() = Some(finished_entry("y", 200)); fx.svc.idle_check().await; assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 1); tokio::time::sleep(std::time::Duration::from_millis(1100)).await; fx.svc.idle_check().await; assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 2); } #[actix_rt::test] async fn busy_queue_blocks_the_free() { let fx = fx_with(None, |c| c.idle_free_secs = 1).await; fx.mock.queue_running.lock().unwrap().push("someone".into()); fx.svc.idle_check().await; tokio::time::sleep(std::time::Duration::from_millis(1100)).await; fx.svc.idle_check().await; assert_eq!(fx.mock.free_calls.load(Ordering::SeqCst), 0); }