//! Integration tests for the Werkbank runner endpoints (WB-05). //! //! Drives the real HTTP handlers (lease/heartbeat/complete) against a live Mongo: //! a runner leases a seeded job, completes it, and the result's findings are //! persisted against the job's target. Also checks the bearer-token gate. Skips //! cleanly when no Mongo is reachable. #![allow(clippy::expect_used, clippy::unwrap_used)] mod common; use std::sync::Arc; use axum::routing::{get, post}; use axum::{middleware, Extension, Router}; use compliance_agent::agent::ComplianceAgent; use compliance_agent::api::handlers::werkbank_jobs; use compliance_agent::database::DatabasePool; use compliance_agent::werkbank::JobQueue; use compliance_core::models::werkbank::{InputRef, Job, JobResult, JobStatus, LeasedJob}; use compliance_core::models::{ Artifact, Finding, OnboardedTarget, PlcFormat, ScanType, Severity, TargetType, }; use common::{dev_config, TEST_RUNNER_TOKEN}; const TENANT: &str = "dev"; /// A running werkbank API on a random port, or `None` if no Mongo. struct Harness { base_url: String, client: reqwest::Client, pool: DatabasePool, db_name: String, } async fn start() -> Option { let uri = std::env::var("TEST_MONGODB_URI") .unwrap_or_else(|_| "mongodb://root:example@localhost:27017/?authSource=admin".into()); let db_name = format!("wba_{}", &uuid::Uuid::new_v4().simple().to_string()[..12]); let pool = match DatabasePool::connect(&uri, &db_name).await { Ok(p) => p, Err(_) => { eprintln!("SKIP werkbank_api: no MongoDB reachable at {uri}"); return None; } }; // Touch the tenant DB so indexes are ensured before the queue is used. pool.for_tenant_id(TENANT).await.expect("tenant db"); let agent = ComplianceAgent::new(dev_config(uri, db_name.clone()), pool.clone()); let app = Router::new() .route("/api/v1/werkbank/jobs/lease", post(werkbank_jobs::lease)) .route( "/api/v1/werkbank/jobs/heartbeat", post(werkbank_jobs::heartbeat), ) .route( "/api/v1/werkbank/jobs/complete", post(werkbank_jobs::complete), ) .route( "/api/v1/werkbank/jobs/enqueue", post(werkbank_jobs::enqueue), ) .route( "/api/v1/werkbank/artifacts/{hash}", get(werkbank_jobs::serve_artifact), ) .layer(middleware::from_fn(werkbank_jobs::require_runner_token)) .layer(Extension(Arc::new(agent))); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); tokio::spawn(async move { axum::serve(listener, app).await.ok(); }); Some(Harness { base_url: format!("http://127.0.0.1:{port}"), client: reqwest::Client::new(), pool, db_name, }) } impl Harness { fn post( &self, path: &str, token: Option<&str>, body: serde_json::Value, ) -> reqwest::RequestBuilder { let mut r = self .client .post(format!("{}{path}", self.base_url)) .json(&body); if let Some(t) = token { r = r.bearer_auth(t); } r } async fn cleanup(&self) { let _ = self .pool .client() .database(&format!("{}_{TENANT}", self.db_name)) .drop() .await; } } fn finding_for(target: &str, fp: &str) -> Finding { let mut f = Finding::new( target.to_string(), fp.to_string(), "ics-probe".to_string(), ScanType::IcsProbe, "Modbus exposed".to_string(), "unauthenticated".to_string(), Severity::Critical, ); f.rule_id = Some("ics-modbus-exposed".to_string()); f } #[tokio::test] async fn lease_complete_persists_findings_against_the_target() { let Some(h) = start().await else { return }; let db = h.pool.for_tenant_id(TENANT).await.unwrap(); let queue = JobQueue::new(&db); // Seed a queued job. let job = Job::plc_provision("job-1", TENANT, "target-1", InputRef::blob("sha256:x"), 180); assert!(queue.enqueue(job, chrono::Utc::now()).await.unwrap()); // Lease it over HTTP. let resp = h .post( "/api/v1/werkbank/jobs/lease", Some(TEST_RUNNER_TOKEN), serde_json::json!({ "tenant": TENANT, "runner_id": "r1", "executor": "docker", "labels": [], "lease_ttl_secs": 60 }), ) .send() .await .unwrap(); assert_eq!(resp.status(), 200, "lease should return a job"); let leased: LeasedJob = resp.json().await.unwrap(); assert_eq!(leased.job.id, "job-1"); // Complete it with a finding. let mut result = JobResult::succeeded("job-1"); result.findings = vec![finding_for("target-1", "fp-abc")]; let resp = h .post( "/api/v1/werkbank/jobs/complete", Some(TEST_RUNNER_TOKEN), serde_json::json!({ "tenant": TENANT, "job_id": "job-1", "lease_token": leased.lease_token, "result": result }), ) .send() .await .unwrap(); assert_eq!(resp.status(), 200); assert!(resp.json::().await.unwrap()["recorded"] .as_bool() .unwrap()); // The job is now succeeded, and the finding was persisted to the target. assert_eq!( queue.get("job-1").await.unwrap().unwrap().status, JobStatus::Succeeded ); let stored = db .findings() .find_one(mongodb::bson::doc! { "fingerprint": "fp-abc" }) .await .unwrap(); assert!(stored.is_some(), "finding should be persisted"); h.cleanup().await; } #[tokio::test] async fn enqueue_extracts_program_stores_a_blob_and_serves_it() { let Some(h) = start().await else { return }; let db = h.pool.for_tenant_id(TENANT).await.unwrap(); // A PlcSps target with a single complete ST program uploaded. let dir = std::env::temp_dir().join(format!("wbq-prog-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let st = dir.join("main.st"); std::fs::write( &st, "PROGRAM Main\nEND_PROGRAM\nCONFIGURATION C\n RESOURCE R\nEND_CONFIGURATION\n", ) .unwrap(); let mut target = OnboardedTarget::new("plc".into(), TargetType::PlcSps); let mut art = Artifact::plc_project("main.st", PlcFormat::StructuredText); art.stored_path = Some(st.to_string_lossy().to_string()); target.artifacts.push(art); let ins = db.onboarded_targets().insert_one(&target).await.unwrap(); let target_id = ins.inserted_id.as_object_id().unwrap().to_hex(); // Enqueue → a plc-provision job whose program is a content-addressed blob. let resp = h .post( "/api/v1/werkbank/jobs/enqueue", Some(TEST_RUNNER_TOKEN), serde_json::json!({ "tenant": TENANT, "target_id": target_id }), ) .send() .await .unwrap(); assert_eq!(resp.status(), 200, "enqueue should succeed"); let body: serde_json::Value = resp.json().await.unwrap(); let job_id = body["job_id"].as_str().unwrap().to_string(); let rec = JobQueue::new(&db).get(&job_id).await.unwrap().unwrap(); let hash = rec .job .inputs .get("program") .and_then(|i| i.blob.clone()) .expect("program blob"); // Serve the blob back and confirm it's the program source (what the runner // would fetch). let served = h .client .get(format!("{}/api/v1/werkbank/artifacts/{hash}", h.base_url)) .bearer_auth(TEST_RUNNER_TOKEN) .send() .await .unwrap(); assert_eq!(served.status(), 200); assert!(served.text().await.unwrap().contains("CONFIGURATION")); h.cleanup().await; let _ = std::fs::remove_dir_all(&dir); } #[tokio::test] async fn empty_queue_leases_nothing() { let Some(h) = start().await else { return }; let resp = h .post( "/api/v1/werkbank/jobs/lease", Some(TEST_RUNNER_TOKEN), serde_json::json!({ "tenant": TENANT, "runner_id": "r1", "executor": "docker", "labels": [], "lease_ttl_secs": 60 }), ) .send() .await .unwrap(); assert_eq!(resp.status(), 204, "no job → 204"); h.cleanup().await; } #[tokio::test] async fn runner_endpoints_require_the_bearer_token() { let Some(h) = start().await else { return }; let body = serde_json::json!({ "tenant": TENANT, "runner_id": "r1", "executor": "docker", "labels": [], "lease_ttl_secs": 60 }); let no_token = h .post("/api/v1/werkbank/jobs/lease", None, body.clone()) .send() .await .unwrap(); assert_eq!(no_token.status(), 401, "missing token → 401"); let bad_token = h .post("/api/v1/werkbank/jobs/lease", Some("wrong"), body) .send() .await .unwrap(); assert_eq!(bad_token.status(), 401, "wrong token → 401"); h.cleanup().await; }