feat(enhance): discard, tests against a mock ComfyUI, docs (#192)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
1 parent
3a21761cfe
commit
97d9301a71
10 files changed
+1165
-48
No files matched your search
@@ -0,0 +1,985 @@
|
||||
//! 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<Vec<Value>>,
|
||||
queue_running: Mutex<Vec<String>>,
|
||||
queue_pending: Mutex<Vec<String>>,
|
||||
/// prompt_id -> history entry
|
||||
history: Mutex<HashMap<String, Value>>,
|
||||
newest: Mutex<Option<Value>>,
|
||||
approve_calls: Mutex<Vec<String>>,
|
||||
approve_path: Mutex<String>,
|
||||
approve_gone: Mutex<bool>,
|
||||
free_calls: AtomicUsize,
|
||||
reject_prompts: Mutex<Option<Value>>,
|
||||
view_404: Mutex<bool>,
|
||||
}
|
||||
|
||||
async fn start_mock() -> (Arc<Mock>, 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<Mock>, body: web::Json<Value>| 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<Mock>| async move {
|
||||
let item = |id: &String| json!([1, id, {}, {}, []]);
|
||||
HttpResponse::Ok().json(json!({
|
||||
"queue_running": m.queue_running.lock().unwrap().iter().map(item).collect::<Vec<_>>(),
|
||||
"queue_pending": m.queue_pending.lock().unwrap().iter().map(item).collect::<Vec<_>>(),
|
||||
}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/history",
|
||||
web::get().to(|m: Data<Mock>| 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<Mock>, id: web::Path<String>| 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<Mock>, q: web::Query<HashMap<String, String>>| 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<Mock>, body: web::Json<Value>| 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<Mock>, body: web::Json<Value>| 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<Mock>,
|
||||
svc: Data<EnhanceService>,
|
||||
state: Data<AppState>,
|
||||
exif: Data<Mutex<Box<dyn ExifDao>>>,
|
||||
conn: Arc<Mutex<SqliteConnection>>,
|
||||
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<String>, 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::<Text, _>(&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<Mutex<Box<dyn EnhanceJobDao>>> = 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<dyn ExifDao>,
|
||||
)),
|
||||
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::<Text, _>(rel)
|
||||
.bind::<Text, _>(&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<Value>) -> (u16, Value, Vec<u8>) {
|
||||
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<Value>) -> (u16, Value, Vec<u8>) {
|
||||
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::<Text, _>(saved_rel)
|
||||
.get_result::<N>(&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);
|
||||
}
|
||||
Reference in new issue
Block a user