compliance-scanner-agent/compliance-agent/tests/common/mod.rs

// Shared test harness for E2E / integration tests.
//
// Spins up the agent API server on a random port with an isolated test
// database. Each test gets a fresh database that is dropped on cleanup.

use std::sync::Arc;

use compliance_agent::agent::ComplianceAgent;
use compliance_agent::api;
use compliance_agent::database::Database;
use compliance_core::AgentConfig;
use secrecy::SecretString;

/// A running test server with a unique database.
pub struct TestServer {
    pub base_url: String,
    pub client: reqwest::Client,
    db_name: String,
    mongodb_uri: String,
}

impl TestServer {
    /// Start an agent API server on a random port with an isolated database.
    pub async fn start() -> Self {
        let mongodb_uri = std::env::var("TEST_MONGODB_URI")
            .unwrap_or_else(|_| "mongodb://root:example@localhost:27017/?authSource=admin".into());

        // Unique database name per test run to avoid collisions
        let db_name = format!("test_{}", uuid::Uuid::new_v4().simple());

        let db = Database::connect(&mongodb_uri, &db_name)
            .await
            .expect("Failed to connect to MongoDB — is it running?");
        db.ensure_indexes().await.expect("Failed to create indexes");

        let config = AgentConfig {
            mongodb_uri: mongodb_uri.clone(),
            mongodb_database: db_name.clone(),
            litellm_url: std::env::var("TEST_LITELLM_URL")
                .unwrap_or_else(|_| "http://localhost:4000".into()),
            litellm_api_key: SecretString::from(String::new()),
            litellm_model: "gpt-4o".into(),
            litellm_embed_model: "text-embedding-3-small".into(),
            agent_port: 0, // not used — we bind ourselves
            scan_schedule: String::new(),
            cve_monitor_schedule: String::new(),
            git_clone_base_path: "/tmp/compliance-scanner-tests/repos".into(),
            ssh_key_path: "/tmp/compliance-scanner-tests/ssh/id_ed25519".into(),
            github_token: None,
            github_webhook_secret: None,
            gitlab_url: None,
            gitlab_token: None,
            gitlab_webhook_secret: None,
            jira_url: None,
            jira_email: None,
            jira_api_token: None,
            jira_project_key: None,
            searxng_url: None,
            nvd_api_key: None,
            keycloak_url: None,
            keycloak_realm: None,
            keycloak_admin_username: None,
            keycloak_admin_password: None,
            pentest_verification_email: None,
            pentest_imap_host: None,
            pentest_imap_port: None,
            pentest_imap_tls: false,
            pentest_imap_username: None,
            pentest_imap_password: None,
        };

        let agent = ComplianceAgent::new(config, db);

        // Build the router with the agent extension
        let app = api::routes::build_router()
            .layer(axum::extract::Extension(Arc::new(agent)))
            .layer(tower_http::cors::CorsLayer::permissive());

        // Bind to port 0 to get a random available port
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
            .await
            .expect("Failed to bind test server");
        let port = listener.local_addr().expect("no local addr").port();

        tokio::spawn(async move {
            axum::serve(listener, app).await.ok();
        });

        let base_url = format!("http://127.0.0.1:{port}");
        let client = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(30))
            .build()
            .expect("Failed to build HTTP client");

        // Wait for server to be ready
        for _ in 0..50 {
            if client
                .get(format!("{base_url}/api/v1/health"))
                .send()
                .await
                .is_ok()
            {
                break;
            }
            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        }

        Self {
            base_url,
            client,
            db_name,
            mongodb_uri,
        }
    }

    /// GET helper
    pub async fn get(&self, path: &str) -> reqwest::Response {
        self.client
            .get(format!("{}{path}", self.base_url))
            .send()
            .await
            .expect("GET request failed")
    }

    /// POST helper with JSON body
    pub async fn post(&self, path: &str, body: &serde_json::Value) -> reqwest::Response {
        self.client
            .post(format!("{}{path}", self.base_url))
            .json(body)
            .send()
            .await
            .expect("POST request failed")
    }

    /// PATCH helper with JSON body
    pub async fn patch(&self, path: &str, body: &serde_json::Value) -> reqwest::Response {
        self.client
            .patch(format!("{}{path}", self.base_url))
            .json(body)
            .send()
            .await
            .expect("PATCH request failed")
    }

    /// DELETE helper
    pub async fn delete(&self, path: &str) -> reqwest::Response {
        self.client
            .delete(format!("{}{path}", self.base_url))
            .send()
            .await
            .expect("DELETE request failed")
    }

    /// Get the unique database name for direct MongoDB access in tests.
    pub fn db_name(&self) -> &str {
        &self.db_name
    }

    /// Drop the test database on cleanup
    pub async fn cleanup(&self) {
        if let Ok(client) = mongodb::Client::with_uri_str(&self.mongodb_uri).await {
            client.database(&self.db_name).drop().await.ok();
        }
    }
}