#[cfg(test)]
mod tests {
use axum::{
body::Body,
http::{Request, StatusCode},
};
use mold_core::{GenerateRequest, GenerateResponse, ImageData};
use mold_inference::progress::ProgressCallback;
use mold_inference::InferenceEngine;
use sha2::{Digest, Sha256};
use std::net::IpAddr;
use std::path::PathBuf;
use std::sync::{
atomic::{AtomicBool, AtomicUsize, Ordering},
Arc, Condvar, Mutex,
};
use std::time::Duration;
use tower::ServiceExt;
use crate::{
routes::create_router,
state::{AppState, EngineSnapshot},
};
fn env_lock() -> &'static std::sync::Mutex<()> {
static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
&ENV_LOCK
}
async fn json_body(resp: axum::http::Response<Body>) -> serde_json::Value {
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
serde_json::from_slice(&body).unwrap()
}
#[derive(Default)]
struct GenerateBlocker {
entered: AtomicBool,
released: Mutex<bool>,
released_cv: Condvar,
}
impl GenerateBlocker {
fn release(&self) {
let mut released = self.released.lock().unwrap();
*released = true;
self.released_cv.notify_all();
}
}
struct MockEngine {
loaded: bool,
fail: bool,
empty_images: bool,
load_count: Arc<AtomicUsize>,
load_delay: Duration,
progress_set_count: Arc<AtomicUsize>,
progress_clear_count: Arc<AtomicUsize>,
generate_blocker: Option<Arc<GenerateBlocker>>,
emit_load_progress: bool,
progress_callback: Option<ProgressCallback>,
}
impl MockEngine {
fn ready() -> Self {
Self {
loaded: true,
fail: false,
empty_images: false,
load_count: Arc::new(AtomicUsize::new(0)),
load_delay: Duration::from_millis(0),
progress_set_count: Arc::new(AtomicUsize::new(0)),
progress_clear_count: Arc::new(AtomicUsize::new(0)),
generate_blocker: None,
emit_load_progress: false,
progress_callback: None,
}
}
fn failing() -> Self {
Self {
loaded: true,
fail: true,
empty_images: false,
load_count: Arc::new(AtomicUsize::new(0)),
load_delay: Duration::from_millis(0),
progress_set_count: Arc::new(AtomicUsize::new(0)),
progress_clear_count: Arc::new(AtomicUsize::new(0)),
generate_blocker: None,
emit_load_progress: false,
progress_callback: None,
}
}
fn empty_images() -> Self {
Self {
loaded: true,
fail: false,
empty_images: true,
load_count: Arc::new(AtomicUsize::new(0)),
load_delay: Duration::from_millis(0),
progress_set_count: Arc::new(AtomicUsize::new(0)),
progress_clear_count: Arc::new(AtomicUsize::new(0)),
generate_blocker: None,
emit_load_progress: false,
progress_callback: None,
}
}
fn unloaded(load_count: Arc<AtomicUsize>, load_delay: Duration) -> Self {
Self {
loaded: false,
fail: false,
empty_images: false,
load_count,
load_delay,
progress_set_count: Arc::new(AtomicUsize::new(0)),
progress_clear_count: Arc::new(AtomicUsize::new(0)),
generate_blocker: None,
emit_load_progress: false,
progress_callback: None,
}
}
fn tracked_progress(
progress_set_count: Arc<AtomicUsize>,
progress_clear_count: Arc<AtomicUsize>,
) -> Self {
Self {
loaded: true,
fail: false,
empty_images: false,
load_count: Arc::new(AtomicUsize::new(0)),
load_delay: Duration::from_millis(0),
progress_set_count,
progress_clear_count,
generate_blocker: None,
emit_load_progress: false,
progress_callback: None,
}
}
fn blocking_generate(blocker: Arc<GenerateBlocker>) -> Self {
Self {
loaded: true,
fail: false,
empty_images: false,
load_count: Arc::new(AtomicUsize::new(0)),
load_delay: Duration::from_millis(0),
progress_set_count: Arc::new(AtomicUsize::new(0)),
progress_clear_count: Arc::new(AtomicUsize::new(0)),
generate_blocker: Some(blocker),
emit_load_progress: false,
progress_callback: None,
}
}
fn unloaded_with_progress() -> Self {
Self {
loaded: false,
fail: false,
empty_images: false,
load_count: Arc::new(AtomicUsize::new(0)),
load_delay: Duration::from_millis(0),
progress_set_count: Arc::new(AtomicUsize::new(0)),
progress_clear_count: Arc::new(AtomicUsize::new(0)),
generate_blocker: None,
emit_load_progress: true,
progress_callback: None,
}
}
}
impl InferenceEngine for MockEngine {
fn generate(&mut self, req: &GenerateRequest) -> anyhow::Result<GenerateResponse> {
if let Some(blocker) = &self.generate_blocker {
blocker.entered.store(true, Ordering::SeqCst);
let released = blocker.released.lock().unwrap();
let _released = blocker
.released_cv
.wait_while(released, |released| !*released)
.unwrap();
}
if self.fail {
anyhow::bail!("mock engine error");
}
let images = if self.empty_images {
vec![]
} else {
vec![ImageData {
data: minimal_png(),
format: req.output_format,
width: req.width,
height: req.height,
index: 0,
}]
};
Ok(GenerateResponse {
images,
generation_time_ms: 1,
model: req.model.clone(),
seed_used: req.seed.unwrap_or(42),
video: None,
})
}
fn model_name(&self) -> &str {
"mock-model"
}
fn is_loaded(&self) -> bool {
self.loaded
}
fn load(&mut self) -> anyhow::Result<()> {
self.load_count.fetch_add(1, Ordering::SeqCst);
if self.emit_load_progress {
if let Some(ref cb) = self.progress_callback {
cb(mold_inference::progress::ProgressEvent::Info {
message: "Converting FP8 checkpoint to Q8 GGUF cache (one-time, may take a few minutes)".to_string(),
});
cb(mold_inference::progress::ProgressEvent::StageStart {
name: "Loading transformer (GPU, quantized)".to_string(),
});
}
}
if !self.load_delay.is_zero() {
std::thread::sleep(self.load_delay);
}
self.loaded = true;
Ok(())
}
fn set_on_progress(&mut self, callback: ProgressCallback) {
self.progress_set_count.fetch_add(1, Ordering::SeqCst);
self.progress_callback = Some(callback);
}
fn clear_on_progress(&mut self) {
self.progress_clear_count.fetch_add(1, Ordering::SeqCst);
self.progress_callback = None;
}
}
fn app_with(engine: MockEngine) -> axum::Router {
let (state, rx) = AppState::with_engine_and_queue(engine);
let worker_state = state.clone();
tokio::spawn(crate::queue::run_queue_worker(rx, worker_state));
create_router(state)
}
fn app_with_state(state: AppState) -> axum::Router {
create_router(state)
}
fn app_empty() -> axum::Router {
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let queue = crate::state::QueueHandle::new(tx);
app_with_state(AppState::empty(mold_core::Config::default(), queue))
}
fn generate_body(prompt: &str, width: u32, height: u32) -> String {
format!(
r#"{{"prompt":"{prompt}","model":"mock-model","width":{width},"height":{height},"steps":4,"batch_size":1,"output_format":"png"}}"#
)
}
fn test_models_dir(name: &str) -> PathBuf {
let unique = format!(
"mold-server-routes-{name}-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
);
std::env::temp_dir().join(unique)
}
fn populate_manifest_files(root: &std::path::Path, model: &str) {
let manifest = mold_core::manifest::find_manifest(model).unwrap();
for file in &manifest.files {
let path = root.join(mold_core::manifest::storage_path(manifest, file));
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent).unwrap();
}
std::fs::write(path, b"test").unwrap();
}
}
fn minimal_png() -> Vec<u8> {
vec![
0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53, 0xDE, 0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, 0x54, 0x08, 0xD7, 0x63, 0xF8, 0xCF, 0xC0, 0x00, 0x00, 0x00, 0x02, 0x00,
0x01, 0xE2, 0x21, 0xBC, 0x33, 0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82, ]
}
#[tokio::test]
async fn health_returns_200() {
let app = app_with(MockEngine::ready());
let resp = app
.oneshot(Request::get("/health").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn health_when_no_model() {
let app = app_empty();
let resp = app
.oneshot(Request::get("/health").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn status_returns_json() {
let app = app_with(MockEngine::ready());
let resp = app
.oneshot(Request::get("/api/status").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let ct = resp
.headers()
.get("content-type")
.unwrap()
.to_str()
.unwrap();
assert!(ct.contains("application/json"));
}
#[tokio::test]
async fn status_when_no_model() {
let app = app_empty();
let resp = app
.oneshot(Request::get("/api/status").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
let status: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(status["models_loaded"], serde_json::json!([]));
assert_eq!(status["busy"], serde_json::json!(false));
assert_eq!(status["current_generation"], serde_json::Value::Null);
}
#[tokio::test]
async fn status_includes_hostname_and_memory_status() {
let app = app_empty();
let resp = app
.oneshot(Request::get("/api/status").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
let status: mold_core::ServerStatus = serde_json::from_slice(&body).unwrap();
assert!(
status.hostname.is_some(),
"server should report its hostname"
);
assert!(
!status.hostname.as_ref().unwrap().is_empty(),
"hostname should not be empty"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn status_does_not_block_during_generation() {
let blocker = Arc::new(GenerateBlocker::default());
let app = app_with(MockEngine::blocking_generate(blocker.clone()));
let generate_task = tokio::spawn({
let app = app.clone();
async move {
app.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(generate_body("a cat", 768, 768)))
.unwrap(),
)
.await
.unwrap()
}
});
tokio::time::timeout(Duration::from_secs(1), async {
while !blocker.entered.load(Ordering::SeqCst) {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("generate should enter the mock engine");
let resp = tokio::time::timeout(
Duration::from_millis(200),
app.clone()
.oneshot(Request::get("/api/status").body(Body::empty()).unwrap()),
)
.await
.expect("/api/status should not block on active generation")
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let status = json_body(resp).await;
assert_eq!(status["busy"], serde_json::json!(true));
assert_eq!(status["current_generation"]["model"], "mock-model");
assert_eq!(
status["current_generation"]["prompt_sha256"],
serde_json::json!(format!("{:x}", Sha256::digest("a cat".as_bytes())))
);
assert!(
status["current_generation"]["started_at_unix_ms"]
.as_u64()
.unwrap()
> 0
);
blocker.release();
let generate_resp = generate_task.await.unwrap();
assert_eq!(generate_resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn list_models_returns_json_array() {
let app = app_with(MockEngine::ready());
let resp = app
.oneshot(Request::get("/api/models").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn list_models_uses_manifest_defaults_for_unpulled() {
let app = app_with(MockEngine::ready());
let resp = app
.oneshot(Request::get("/api/models").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
let models: Vec<serde_json::Value> = serde_json::from_slice(&body).unwrap();
let sd15 = models.iter().find(|m| m["name"] == "sd15:fp16");
if let Some(sd15) = sd15 {
assert_eq!(sd15["default_width"], 512, "SD1.5 width should be 512");
assert_eq!(sd15["default_height"], 512, "SD1.5 height should be 512");
assert_eq!(sd15["default_steps"], 25, "SD1.5 steps should be 25");
assert_eq!(
sd15["default_guidance"], 7.5,
"SD1.5 guidance should be 7.5"
);
}
let schnell = models.iter().find(|m| m["name"] == "flux-schnell:q8");
if let Some(schnell) = schnell {
assert_eq!(schnell["default_width"], 1024);
assert_eq!(schnell["default_height"], 1024);
assert_eq!(schnell["default_steps"], 4);
}
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn list_models_reports_server_disk_and_remaining_download_bytes() {
let _lock = env_lock().lock().unwrap_or_else(|e| e.into_inner());
let models_dir = test_models_dir("remote-catalog");
populate_manifest_files(&models_dir, "flux-schnell:q8");
std::env::set_var("MOLD_MODELS_DIR", &models_dir);
let app = app_empty();
let resp = app
.oneshot(Request::get("/api/models").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
let models: Vec<serde_json::Value> = serde_json::from_slice(&body).unwrap();
let downloaded = models
.iter()
.find(|m| m["name"] == "flux-schnell:q8")
.expect("flux-schnell:q8 should be present");
assert_eq!(downloaded["downloaded"], serde_json::json!(true));
assert!(
downloaded["remaining_download_bytes"].is_number(),
"downloaded model should expose remaining download bytes"
);
assert!(
downloaded["disk_usage_bytes"].as_u64().unwrap() > 0,
"downloaded model should report server disk usage"
);
let available = models
.iter()
.find(|m| m["name"] == "flux-dev:q8")
.expect("flux-dev:q8 should be present");
assert_eq!(available["downloaded"], serde_json::json!(false));
assert!(
available["remaining_download_bytes"].is_number(),
"available model should expose server-side remaining bytes even when fully cached"
);
std::env::remove_var("MOLD_MODELS_DIR");
let _ = std::fs::remove_dir_all(models_dir);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn list_models_does_not_block_during_generation() {
let blocker = Arc::new(GenerateBlocker::default());
let app = app_with(MockEngine::blocking_generate(blocker.clone()));
let generate_task = tokio::spawn({
let app = app.clone();
async move {
app.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(generate_body("a cat", 768, 768)))
.unwrap(),
)
.await
.unwrap()
}
});
tokio::time::timeout(Duration::from_secs(1), async {
while !blocker.entered.load(Ordering::SeqCst) {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("generate should enter the mock engine");
let resp = tokio::time::timeout(
Duration::from_millis(200),
app.clone()
.oneshot(Request::get("/api/models").body(Body::empty()).unwrap()),
)
.await
.expect("/api/models should not block on active generation")
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
blocker.release();
let generate_resp = generate_task.await.unwrap();
assert_eq!(generate_resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn generate_empty_prompt_returns_422() {
let app = app_with(MockEngine::ready());
let body = generate_body("", 768, 768);
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::UNPROCESSABLE_ENTITY);
let body = json_body(resp).await;
assert_eq!(body["code"], "VALIDATION_ERROR");
assert!(body["error"].as_str().unwrap().contains("prompt"));
}
#[tokio::test]
async fn generate_zero_width_returns_422() {
let app = app_with(MockEngine::ready());
let body = generate_body("a cat", 0, 768);
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::UNPROCESSABLE_ENTITY);
let body = json_body(resp).await;
assert_eq!(body["code"], "VALIDATION_ERROR");
}
#[tokio::test]
async fn generate_non_multiple_of_16_returns_422() {
let app = app_with(MockEngine::ready());
let body = generate_body("a cat", 769, 768);
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::UNPROCESSABLE_ENTITY);
let body = json_body(resp).await;
assert_eq!(body["code"], "VALIDATION_ERROR");
}
#[tokio::test]
async fn generate_oversized_returns_422() {
let app = app_with(MockEngine::ready());
let body = generate_body("a cat", 1408, 1408);
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::UNPROCESSABLE_ENTITY);
let body = json_body(resp).await;
assert_eq!(body["code"], "VALIDATION_ERROR");
}
#[tokio::test]
async fn generate_zero_steps_returns_422() {
let app = app_with(MockEngine::ready());
let body = r#"{"prompt":"a cat","model":"mock-model","width":768,"height":768,"steps":0,"batch_size":1,"output_format":"png"}"#;
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::UNPROCESSABLE_ENTITY);
let body = json_body(resp).await;
assert_eq!(body["code"], "VALIDATION_ERROR");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn generate_valid_request_returns_image_bytes() {
let app = app_with(MockEngine::ready());
let body = generate_body("a glowing robot", 768, 768);
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let ct = resp
.headers()
.get("content-type")
.unwrap()
.to_str()
.unwrap();
assert_eq!(ct, "image/png");
assert!(
resp.headers().contains_key("x-mold-seed-used"),
"response should include x-mold-seed-used header"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn generate_engine_error_returns_500() {
let app = app_with(MockEngine::failing());
let body = generate_body("a cat", 768, 768);
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
let body = json_body(resp).await;
assert_eq!(body["code"], "INFERENCE_ERROR");
assert!(body["error"]
.as_str()
.unwrap()
.contains("mock engine error"));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn generate_empty_images_returns_500() {
let app = app_with(MockEngine::empty_images());
let body = generate_body("a cat", 768, 768);
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
let body = json_body(resp).await;
assert_eq!(body["code"], "INFERENCE_ERROR");
assert!(body["error"]
.as_str()
.unwrap()
.contains("returned no images"));
}
#[tokio::test]
async fn generate_unknown_model_returns_400() {
let app = app_empty();
let body = r#"{"prompt":"a cat","model":"nonexistent-model-xyz","width":768,"height":768,"steps":4,"batch_size":1,"output_format":"png"}"#;
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let body = json_body(resp).await;
assert_eq!(body["code"], "UNKNOWN_MODEL");
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn generate_known_model_not_downloaded_returns_404() {
let _lock = env_lock().lock().unwrap_or_else(|e| e.into_inner());
let models_dir = test_models_dir("generate-not-downloaded");
std::fs::create_dir_all(&models_dir).unwrap();
std::env::set_var("MOLD_MODELS_DIR", &models_dir);
let app = app_empty();
let body = r#"{"prompt":"a cat","model":"flux-schnell:q8","width":768,"height":768,"steps":4,"batch_size":1,"output_format":"png"}"#;
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::NOT_FOUND);
let body = json_body(resp).await;
assert_eq!(body["code"], "MODEL_NOT_FOUND");
std::env::remove_var("MOLD_MODELS_DIR");
let _ = std::fs::remove_dir_all(models_dir);
}
#[tokio::test]
async fn openapi_json_returns_valid_spec() {
let app = app_with(MockEngine::ready());
let resp = app
.oneshot(
Request::get("/api/openapi.json")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
let spec: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert!(
spec["openapi"].is_string(),
"spec should have openapi version"
);
assert!(spec["paths"].is_object(), "spec should have paths");
assert!(
spec["paths"]["/api/generate"].is_object(),
"spec should have /api/generate path"
);
}
#[tokio::test]
async fn docs_returns_html() {
let app = app_with(MockEngine::ready());
let resp = app
.oneshot(Request::get("/api/docs").body(Body::empty()).unwrap())
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let ct = resp
.headers()
.get("content-type")
.unwrap()
.to_str()
.unwrap();
assert!(
ct.contains("text/html"),
"docs should return HTML, got: {ct}"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stream_valid_request_returns_sse() {
let app = app_with(MockEngine::ready());
let body = generate_body("a robot", 768, 768);
let resp = app
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let ct = resp
.headers()
.get("content-type")
.unwrap()
.to_str()
.unwrap();
assert!(
ct.contains("text/event-stream"),
"stream should return text/event-stream, got: {ct}"
);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
let text = String::from_utf8_lossy(&body);
assert!(
text.contains("event: complete"),
"stream should contain a complete event"
);
assert!(
text.contains("\"image\""),
"complete event should contain base64 image"
);
}
#[tokio::test]
async fn stream_empty_prompt_returns_422() {
let app = app_with(MockEngine::ready());
let body = generate_body("", 768, 768);
let resp = app
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::UNPROCESSABLE_ENTITY);
let body = json_body(resp).await;
assert_eq!(body["code"], "VALIDATION_ERROR");
}
#[tokio::test]
async fn stream_unknown_model_returns_400() {
let app = app_empty();
let body = r#"{"prompt":"a cat","model":"nonexistent-model-xyz","width":768,"height":768,"steps":4,"batch_size":1,"output_format":"png"}"#;
let resp = app
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
let body = json_body(resp).await;
assert_eq!(body["code"], "UNKNOWN_MODEL");
}
#[tokio::test]
#[allow(clippy::await_holding_lock)]
async fn stream_known_model_not_downloaded_returns_404() {
let _lock = env_lock().lock().unwrap_or_else(|e| e.into_inner());
let models_dir = test_models_dir("stream-not-downloaded");
std::fs::create_dir_all(&models_dir).unwrap();
std::env::set_var("MOLD_MODELS_DIR", &models_dir);
let app = app_empty();
let body = r#"{"prompt":"a cat","model":"flux-schnell:q8","width":768,"height":768,"steps":4,"batch_size":1,"output_format":"png"}"#;
let resp = app
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::NOT_FOUND);
let body = json_body(resp).await;
assert_eq!(body["code"], "MODEL_NOT_FOUND");
std::env::remove_var("MOLD_MODELS_DIR");
let _ = std::fs::remove_dir_all(models_dir);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stream_engine_error_returns_sse_error() {
let app = app_with(MockEngine::failing());
let body = generate_body("a cat", 768, 768);
let resp = app
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
let text = String::from_utf8_lossy(&body);
assert!(
text.contains("event: error"),
"stream should contain an error event"
);
assert!(
text.contains("mock engine error"),
"error event should contain the engine error message"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stream_empty_images_returns_sse_error() {
let app = app_with(MockEngine::empty_images());
let body = generate_body("a cat", 768, 768);
let resp = app
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(body))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
let text = String::from_utf8_lossy(&body);
assert!(text.contains("event: error"));
assert!(text.contains("returned no images"));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn reused_engine_clears_progress_callbacks_between_stream_and_generate() {
let progress_set_count = Arc::new(AtomicUsize::new(0));
let progress_clear_count = Arc::new(AtomicUsize::new(0));
let app = app_with(MockEngine::tracked_progress(
progress_set_count.clone(),
progress_clear_count.clone(),
));
let stream_resp = app
.clone()
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(generate_body("a robot", 768, 768)))
.unwrap(),
)
.await
.unwrap();
assert_eq!(stream_resp.status(), StatusCode::OK);
let _ = axum::body::to_bytes(stream_resp.into_body(), 1024 * 1024)
.await
.unwrap();
assert_eq!(progress_set_count.load(Ordering::SeqCst), 2);
assert_eq!(progress_clear_count.load(Ordering::SeqCst), 1);
let generate_resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(generate_body("a robot", 768, 768)))
.unwrap(),
)
.await
.unwrap();
assert_eq!(generate_resp.status(), StatusCode::OK);
assert_eq!(progress_set_count.load(Ordering::SeqCst), 2);
assert_eq!(progress_clear_count.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn unload_loaded_model_returns_200() {
let app = app_with(MockEngine::ready());
let resp = app
.oneshot(
Request::delete("/api/models/unload")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
assert!(String::from_utf8_lossy(&body).contains("unloaded mock-model"));
}
#[tokio::test]
async fn unload_drops_engine_entirely() {
let state = AppState::with_engine(MockEngine::ready());
let app = app_with_state(state.clone());
let resp = app
.oneshot(
Request::delete("/api/models/unload")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let cache = state.model_cache.lock().await;
assert!(
cache.active_model().is_none(),
"no model should be active after unload"
);
}
#[tokio::test]
async fn unload_clears_snapshot_model_name() {
let state = AppState::with_engine(MockEngine::ready());
let app = app_with_state(state.clone());
{
let snapshot = state.engine_snapshot.read().await;
assert!(snapshot.model_name.is_some());
assert!(snapshot.is_loaded);
}
let resp = app
.oneshot(
Request::delete("/api/models/unload")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let snapshot = state.engine_snapshot.read().await;
assert!(
snapshot.model_name.is_none(),
"snapshot model_name should be None after unload"
);
assert!(!snapshot.is_loaded);
}
#[tokio::test]
async fn unload_no_model_returns_200_with_message() {
let (tx, _rx) = tokio::sync::mpsc::channel(16);
let queue = crate::state::QueueHandle::new(tx);
let app = app_with_state(AppState::empty(mold_core::Config::default(), queue));
let resp = app
.oneshot(
Request::delete("/api/models/unload")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
assert!(String::from_utf8_lossy(&body).contains("no model loaded"));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn concurrent_requests_only_load_existing_engine_once() {
let load_count = Arc::new(AtomicUsize::new(0));
let (tx, rx) = tokio::sync::mpsc::channel(16);
let queue = crate::state::QueueHandle::new(tx);
let engine = MockEngine::unloaded(load_count.clone(), Duration::from_millis(50));
let mut cache = crate::model_cache::ModelCache::new(3);
cache.insert(Box::new(engine), 0);
let state = AppState {
model_cache: Arc::new(tokio::sync::Mutex::new(cache)),
engine_snapshot: Arc::new(tokio::sync::RwLock::new(EngineSnapshot {
model_name: Some("mock-model".to_string()),
is_loaded: false,
cached_models: vec!["mock-model".to_string()],
})),
active_generation: Arc::new(std::sync::RwLock::new(None)),
config: Arc::new(tokio::sync::RwLock::new(mold_core::Config::default())),
start_time: std::time::Instant::now(),
model_load_lock: Arc::new(tokio::sync::Mutex::new(())),
pull_lock: Arc::new(tokio::sync::Mutex::new(())),
queue,
shared_pool: Arc::new(std::sync::Mutex::new(
mold_inference::shared_pool::SharedPool::new(),
)),
shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
};
let worker_state = state.clone();
tokio::spawn(crate::queue::run_queue_worker(rx, worker_state));
let app = app_with_state(state);
let req1 = Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(generate_body("a cat", 768, 768)))
.unwrap();
let req2 = Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(generate_body("a cat", 768, 768)))
.unwrap();
let (resp1, resp2) = tokio::join!(app.clone().oneshot(req1), app.oneshot(req2));
assert_eq!(resp1.unwrap().status(), StatusCode::OK);
assert_eq!(resp2.unwrap().status(), StatusCode::OK);
assert_eq!(load_count.load(Ordering::SeqCst), 1);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stream_delivers_load_progress_events() {
let (tx, rx) = tokio::sync::mpsc::channel(16);
let queue = crate::state::QueueHandle::new(tx);
let engine = MockEngine::unloaded_with_progress();
let mut cache = crate::model_cache::ModelCache::new(3);
cache.insert(Box::new(engine), 0);
let state = AppState {
model_cache: Arc::new(tokio::sync::Mutex::new(cache)),
engine_snapshot: Arc::new(tokio::sync::RwLock::new(EngineSnapshot {
model_name: Some("mock-model".to_string()),
is_loaded: false,
cached_models: vec!["mock-model".to_string()],
})),
active_generation: Arc::new(std::sync::RwLock::new(None)),
config: Arc::new(tokio::sync::RwLock::new(mold_core::Config::default())),
start_time: std::time::Instant::now(),
model_load_lock: Arc::new(tokio::sync::Mutex::new(())),
pull_lock: Arc::new(tokio::sync::Mutex::new(())),
queue,
shared_pool: Arc::new(std::sync::Mutex::new(
mold_inference::shared_pool::SharedPool::new(),
)),
shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
};
let worker_state = state.clone();
tokio::spawn(crate::queue::run_queue_worker(rx, worker_state));
let app = app_with_state(state);
let resp = app
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(generate_body("a cat", 768, 768)))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
let text = String::from_utf8_lossy(&body);
assert!(
text.contains("Converting FP8 checkpoint"),
"SSE stream should contain FP8 conversion progress info event, got: {text}"
);
assert!(
text.contains("Loading transformer"),
"SSE stream should contain model loading stage event, got: {text}"
);
assert!(
text.contains("event: complete"),
"SSE stream should contain complete event, got: {text}"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn concurrent_stream_requests_both_complete() {
let blocker = Arc::new(GenerateBlocker::default());
let (state, rx) =
AppState::with_engine_and_queue(MockEngine::blocking_generate(blocker.clone()));
let worker_state = state.clone();
tokio::spawn(crate::queue::run_queue_worker(rx, worker_state));
let app = app_with_state(state);
let resp1_future = {
let app = app.clone();
tokio::spawn(async move {
let resp = app
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(generate_body("request one", 768, 768)))
.unwrap(),
)
.await
.unwrap();
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
String::from_utf8_lossy(&body).to_string()
})
};
tokio::time::sleep(Duration::from_millis(50)).await;
let resp2_future = {
let app = app.clone();
tokio::spawn(async move {
let resp = app
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(generate_body("request two", 768, 768)))
.unwrap(),
)
.await
.unwrap();
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
String::from_utf8_lossy(&body).to_string()
})
};
tokio::time::sleep(Duration::from_millis(50)).await;
blocker.release();
let text1 = resp1_future.await.unwrap();
let text2 = resp2_future.await.unwrap();
assert!(
text1.contains("event: complete"),
"first request should complete, got: {text1}"
);
assert!(
text2.contains("event: complete"),
"second request should complete, got: {text2}"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn queued_stream_receives_position_event() {
let (state, rx) = AppState::with_engine_and_queue(MockEngine::ready());
let queue = state.queue.clone();
let worker_state = state.clone();
let app = app_with_state(state);
let _resp1 = {
let app = app.clone();
tokio::spawn(async move {
app.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(generate_body("first", 768, 768)))
.unwrap(),
)
.await
})
};
while queue.pending() < 1 {
tokio::time::sleep(Duration::from_millis(1)).await;
}
let resp2 = {
let app = app.clone();
tokio::spawn(async move {
let resp = app
.oneshot(
Request::post("/api/generate/stream")
.header("content-type", "application/json")
.body(Body::from(generate_body("second", 768, 768)))
.unwrap(),
)
.await
.unwrap();
let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024)
.await
.unwrap();
String::from_utf8_lossy(&body).to_string()
})
};
while queue.pending() < 2 {
tokio::time::sleep(Duration::from_millis(1)).await;
}
tokio::spawn(crate::queue::run_queue_worker(rx, worker_state));
let text2 = resp2.await.unwrap();
assert!(
text2.contains(r#""type":"queued""#),
"second request should receive a queued event, got: {text2}"
);
assert!(
text2.contains(r#""position":1"#),
"second request should be at position 1, got: {text2}"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn non_streaming_generate_queues_correctly() {
let app = app_with(MockEngine::ready());
let resp1 = {
let app = app.clone();
tokio::spawn(async move {
app.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(generate_body("request one", 768, 768)))
.unwrap(),
)
.await
.unwrap()
})
};
let resp2 = {
let app = app.clone();
tokio::spawn(async move {
app.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(generate_body("request two", 768, 768)))
.unwrap(),
)
.await
.unwrap()
})
};
let (r1, r2) = tokio::join!(resp1, resp2);
assert_eq!(r1.unwrap().status(), StatusCode::OK);
assert_eq!(r2.unwrap().status(), StatusCode::OK);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn snapshot_consistent_after_queue_load() {
let load_count = Arc::new(AtomicUsize::new(0));
let (tx, rx) = tokio::sync::mpsc::channel(16);
let queue = crate::state::QueueHandle::new(tx);
let engine = MockEngine::unloaded(load_count, Duration::from_millis(10));
let mut cache = crate::model_cache::ModelCache::new(3);
cache.insert(Box::new(engine), 0);
let state = AppState {
model_cache: Arc::new(tokio::sync::Mutex::new(cache)),
engine_snapshot: Arc::new(tokio::sync::RwLock::new(EngineSnapshot {
model_name: Some("mock-model".to_string()),
is_loaded: false,
cached_models: vec!["mock-model".to_string()],
})),
active_generation: Arc::new(std::sync::RwLock::new(None)),
config: Arc::new(tokio::sync::RwLock::new(mold_core::Config::default())),
start_time: std::time::Instant::now(),
model_load_lock: Arc::new(tokio::sync::Mutex::new(())),
pull_lock: Arc::new(tokio::sync::Mutex::new(())),
queue,
shared_pool: Arc::new(std::sync::Mutex::new(
mold_inference::shared_pool::SharedPool::new(),
)),
shutdown_tx: Arc::new(tokio::sync::Mutex::new(None)),
upscaler_cache: Arc::new(std::sync::Mutex::new(None)),
};
let worker_state = state.clone();
tokio::spawn(crate::queue::run_queue_worker(rx, worker_state));
let app = app_with_state(state.clone());
let resp = app
.oneshot(
Request::post("/api/generate")
.header("content-type", "application/json")
.body(Body::from(generate_body("a cat", 768, 768)))
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let snapshot = state.engine_snapshot.read().await;
assert_eq!(
snapshot.model_name.as_deref(),
Some("mock-model"),
"snapshot should reflect the loaded model"
);
assert!(snapshot.is_loaded, "snapshot should show model as loaded");
}
fn app_with_auth(auth_state: crate::auth::AuthState) -> axum::Router {
let app = app_empty();
app.layer(axum::middleware::from_fn(crate::auth::require_api_key))
.layer(axum::middleware::from_fn_with_state(
auth_state,
crate::auth::inject_auth_state,
))
}
#[tokio::test]
async fn auth_rejects_missing_api_key() {
let keys = std::collections::HashSet::from(["test-key".to_string()]);
let auth = Some(std::sync::Arc::new(crate::auth::ApiKeySet::new(keys)));
let app = app_with_auth(auth);
let resp = app
.oneshot(
Request::builder()
.uri("/api/status")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
let body = json_body(resp).await;
assert_eq!(body["code"], "UNAUTHORIZED");
}
#[tokio::test]
async fn auth_rejects_invalid_api_key() {
let keys = std::collections::HashSet::from(["test-key".to_string()]);
let auth = Some(std::sync::Arc::new(crate::auth::ApiKeySet::new(keys)));
let app = app_with_auth(auth);
let resp = app
.oneshot(
Request::builder()
.uri("/api/status")
.header("x-api-key", "wrong-key")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn auth_allows_valid_api_key() {
let keys = std::collections::HashSet::from(["test-key".to_string()]);
let auth = Some(std::sync::Arc::new(crate::auth::ApiKeySet::new(keys)));
let app = app_with_auth(auth);
let resp = app
.oneshot(
Request::builder()
.uri("/api/status")
.header("x-api-key", "test-key")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn auth_health_exempt() {
let keys = std::collections::HashSet::from(["test-key".to_string()]);
let auth = Some(std::sync::Arc::new(crate::auth::ApiKeySet::new(keys)));
let app = app_with_auth(auth);
let resp = app
.oneshot(
Request::builder()
.uri("/health")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn auth_docs_exempt() {
let keys = std::collections::HashSet::from(["test-key".to_string()]);
let auth = Some(std::sync::Arc::new(crate::auth::ApiKeySet::new(keys)));
let app = app_with_auth(auth);
let resp = app
.oneshot(
Request::builder()
.uri("/api/docs")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn auth_openapi_exempt() {
let keys = std::collections::HashSet::from(["test-key".to_string()]);
let auth = Some(std::sync::Arc::new(crate::auth::ApiKeySet::new(keys)));
let app = app_with_auth(auth);
let resp = app
.oneshot(
Request::builder()
.uri("/api/openapi.json")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn auth_disabled_when_none() {
let app = app_with_auth(None);
let resp = app
.oneshot(
Request::builder()
.uri("/api/status")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn auth_supports_multiple_keys() {
let keys =
std::collections::HashSet::from(["key-alpha".to_string(), "key-beta".to_string()]);
let auth = Some(std::sync::Arc::new(crate::auth::ApiKeySet::new(keys)));
let app = app_with_auth(auth);
let resp = app
.clone()
.oneshot(
Request::builder()
.uri("/api/status")
.header("x-api-key", "key-alpha")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let resp = app
.oneshot(
Request::builder()
.uri("/api/status")
.header("x-api-key", "key-beta")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
#[tokio::test]
async fn request_id_generated() {
let app = app_empty().layer(axum::middleware::from_fn(
crate::request_id::request_id_middleware,
));
let resp = app
.oneshot(
Request::builder()
.uri("/health")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert!(resp.headers().contains_key("x-request-id"));
}
#[tokio::test]
async fn request_id_preserved() {
let app = app_empty().layer(axum::middleware::from_fn(
crate::request_id::request_id_middleware,
));
let resp = app
.oneshot(
Request::builder()
.uri("/health")
.header("x-request-id", "my-id-123")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(
resp.headers()
.get("x-request-id")
.unwrap()
.to_str()
.unwrap(),
"my-id-123"
);
}
#[test]
fn rate_limit_parse_specs() {
use crate::rate_limit::RouteTier;
use axum::http::Method;
assert_eq!(
crate::rate_limit::classify_route("/api/generate", &Method::POST),
Some(RouteTier::Generation)
);
assert_eq!(
crate::rate_limit::classify_route("/api/generate/stream", &Method::POST),
Some(RouteTier::Generation)
);
assert_eq!(
crate::rate_limit::classify_route("/api/status", &Method::GET),
Some(RouteTier::Read)
);
assert_eq!(
crate::rate_limit::classify_route("/health", &Method::GET),
None
);
}
#[tokio::test]
async fn auth_enforced_on_unmatched_404_paths() {
let keys = std::collections::HashSet::from(["test-key".to_string()]);
let auth = Some(std::sync::Arc::new(crate::auth::ApiKeySet::new(keys)));
let app = app_with_auth(auth);
let resp = app
.oneshot(
Request::builder()
.uri("/nonexistent")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(
resp.status(),
StatusCode::UNAUTHORIZED,
"unmatched paths must still require auth"
);
}
#[test]
fn rate_limiter_map_bounded() {
use crate::rate_limit::MAX_LIMITER_ENTRIES;
let quota = governor::Quota::per_second(std::num::NonZeroU32::new(10).unwrap())
.allow_burst(std::num::NonZeroU32::new(10).unwrap());
let state = crate::rate_limit::RateLimitState::new(quota, quota);
for i in 0..MAX_LIMITER_ENTRIES {
let ip = IpAddr::V4(std::net::Ipv4Addr::from((i as u32).to_be_bytes()));
state.get_generation_limiter(ip);
}
let ip = IpAddr::V4(std::net::Ipv4Addr::new(255, 255, 255, 255));
state.get_generation_limiter(ip);
let map = state.generation_limiters.lock().unwrap();
assert!(map.len() <= 1, "map should be evicted, got {}", map.len());
}
}