compliance-scanner-agent/compliance-agent/tests/werkbank_api.rs

//! 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<Harness> {
    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::<serde_json::Value>().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;
}