photoncloud-monorepo/plasmavmc/crates/plasmavmc-server/tests/creditservice_integration.rs

//! CreditService integration test for PlasmaVMC
//!
//! Tests the 2-phase admission control flow:
//! 1. check_quota - validates balance/quota limits
//! 2. reserve_credits - reserves credits with TTL (Phase 1)
//! 3. [Create Resource] - actual VM creation
//! 4. commit_reservation - commits credits on success (Phase 2)
//! 5. release_reservation - releases credits on failure (rollback)

use creditservice_api::{CreditServiceImpl, CreditStorage, InMemoryStorage};
use creditservice_client::Client as CreditServiceClient;
use creditservice_proto::credit_service_server::CreditServiceServer;
use plasmavmc_api::proto::{
    vm_service_client::VmServiceClient, CreateVmRequest, DeleteVmRequest,
    HypervisorType as ProtoHypervisorType, VmSpec,
};
use plasmavmc_hypervisor::HypervisorRegistry;
use plasmavmc_kvm::KvmBackend;
use plasmavmc_server::VmServiceImpl;
use std::sync::Arc;
use tonic::transport::{Channel, Server};
use tonic::codegen::InterceptedService;
use tonic::service::Interceptor;
use tonic::Request;

struct OrgProjectInterceptor {
    org: String,
    project: String,
}

impl Interceptor for OrgProjectInterceptor {
    fn call(&mut self, mut req: Request<()>) -> Result<Request<()>, tonic::Status> {
        req.metadata_mut().insert("org-id", self.org.parse().unwrap());
        req.metadata_mut().insert("project-id", self.project.parse().unwrap());
        Ok(req)
    }
}

async fn vm_client_with_meta(addr: &str, org: &str, project: &str) -> VmServiceClient<InterceptedService<Channel, OrgProjectInterceptor>> {
    let channel = Channel::from_shared(format!("http://{}", addr)).unwrap().connect().await.unwrap();
    VmServiceClient::with_interceptor(channel, OrgProjectInterceptor { org: org.to_string(), project: project.to_string() })
}

/// Test that CreditService admission control denies VM creation when quota/balance insufficient
#[tokio::test]
#[ignore = "requires PLASMAVMC_QEMU_PATH and PLASMAVMC_QCOW2_PATH"]
async fn creditservice_admission_control_deny() {
    // Skip if QEMU not available
    let qemu = std::env::var("PLASMAVMC_QEMU_PATH").unwrap_or_else(|_| "/usr/bin/qemu-system-x86_64".into());
    if !std::path::Path::new(&qemu).exists() {
        eprintln!("Skipping: QEMU not available at {}", qemu);
        return;
    }

    // 1. Start CreditService
    let credit_addr = "127.0.0.1:50090";
    let storage: Arc<dyn CreditStorage> = InMemoryStorage::new();
    let credit_svc = CreditServiceImpl::new(storage.clone());

    tokio::spawn(async move {
        Server::builder()
            .add_service(CreditServiceServer::new(credit_svc))
            .serve(credit_addr.parse().unwrap())
            .await
            .unwrap();
    });
    tokio::time::sleep(std::time::Duration::from_millis(100)).await;

    // 2. Create wallet with ZERO balance (should deny all requests)
    let mut credit_client = CreditServiceClient::connect(format!("http://{}", credit_addr)).await.unwrap();
    let _wallet = credit_client.create_wallet("proj1", "org1", 0).await.unwrap();

    // 3. Set CREDITSERVICE_ENDPOINT for PlasmaVMC to connect
    std::env::set_var("CREDITSERVICE_ENDPOINT", format!("http://{}", credit_addr));

    // 4. Start PlasmaVMC
    let plasmavmc_addr = "127.0.0.1:50091";
    let registry = Arc::new(HypervisorRegistry::new());
    registry.register(Arc::new(KvmBackend::with_defaults()));
    let vm_svc = VmServiceImpl::new(registry).await.unwrap();

    tokio::spawn(async move {
        Server::builder()
            .add_service(plasmavmc_api::proto::vm_service_server::VmServiceServer::new(vm_svc))
            .serve(plasmavmc_addr.parse().unwrap())
            .await
            .unwrap();
    });
    tokio::time::sleep(std::time::Duration::from_millis(200)).await;

    // 5. Try to create VM - should fail with resource_exhausted
    let mut vm_client = vm_client_with_meta(plasmavmc_addr, "org1", "proj1").await;

    let result = vm_client.create_vm(CreateVmRequest {
        name: "test-vm".into(),
        org_id: "org1".into(),
        project_id: "proj1".into(),
        spec: Some(VmSpec {
            cpu: Some(plasmavmc_api::proto::CpuSpec {
                vcpus: 2,
                cores_per_socket: 1,
                sockets: 1,
                cpu_model: String::new(),
            }),
            memory: Some(plasmavmc_api::proto::MemorySpec {
                size_mib: 1024,
                hugepages: false,
            }),
            disks: vec![],
            network: vec![],
            boot: None,
            security: None,
        }),
        hypervisor: ProtoHypervisorType::Kvm as i32,
        metadata: Default::default(),
        labels: Default::default(),
    }).await;

    // Should fail with resource_exhausted (insufficient balance)
    assert!(result.is_err());
    let err = result.unwrap_err();
    assert_eq!(err.code(), tonic::Code::ResourceExhausted, "Expected ResourceExhausted, got: {:?}", err);

    // Clean up
    std::env::remove_var("CREDITSERVICE_ENDPOINT");
}

/// Test that CreditService admission control allows VM creation with sufficient balance
/// and properly commits/releases credits
#[tokio::test]
#[ignore = "requires PLASMAVMC_QEMU_PATH and PLASMAVMC_QCOW2_PATH"]
async fn creditservice_admission_control_allow() {
    // Skip if QEMU not available
    let qemu = std::env::var("PLASMAVMC_QEMU_PATH").unwrap_or_else(|_| "/usr/bin/qemu-system-x86_64".into());
    let qcow = match std::env::var("PLASMAVMC_QCOW2_PATH") {
        Ok(path) => path,
        Err(_) => {
            eprintln!("Skipping: PLASMAVMC_QCOW2_PATH not set");
            return;
        }
    };
    if !std::path::Path::new(&qemu).exists() || !std::path::Path::new(&qcow).exists() {
        eprintln!("Skipping: QEMU or qcow2 not available");
        return;
    }

    // 1. Start CreditService
    let credit_addr = "127.0.0.1:50092";
    let storage: Arc<dyn CreditStorage> = InMemoryStorage::new();
    let credit_svc = CreditServiceImpl::new(storage.clone());

    tokio::spawn(async move {
        Server::builder()
            .add_service(CreditServiceServer::new(credit_svc))
            .serve(credit_addr.parse().unwrap())
            .await
            .unwrap();
    });
    tokio::time::sleep(std::time::Duration::from_millis(100)).await;

    // 2. Create wallet with sufficient balance
    // Cost = vcpus * 10 + memory_gb * 5 = 2 * 10 + 1 * 5 = 25
    let mut credit_client = CreditServiceClient::connect(format!("http://{}", credit_addr)).await.unwrap();
    let wallet = credit_client.create_wallet("proj2", "org2", 1000).await.unwrap();
    assert_eq!(wallet.balance, 1000);

    // 3. Set CREDITSERVICE_ENDPOINT for PlasmaVMC to connect
    std::env::set_var("CREDITSERVICE_ENDPOINT", format!("http://{}", credit_addr));

    // 4. Start PlasmaVMC
    let plasmavmc_addr = "127.0.0.1:50093";
    let registry = Arc::new(HypervisorRegistry::new());
    registry.register(Arc::new(KvmBackend::with_defaults()));
    let vm_svc = VmServiceImpl::new(registry).await.unwrap();

    tokio::spawn(async move {
        Server::builder()
            .add_service(plasmavmc_api::proto::vm_service_server::VmServiceServer::new(vm_svc))
            .serve(plasmavmc_addr.parse().unwrap())
            .await
            .unwrap();
    });
    tokio::time::sleep(std::time::Duration::from_millis(200)).await;

    // 5. Create VM - should succeed
    let mut vm_client = vm_client_with_meta(plasmavmc_addr, "org2", "proj2").await;

    let vm = vm_client.create_vm(CreateVmRequest {
        name: "test-vm-allowed".into(),
        org_id: "org2".into(),
        project_id: "proj2".into(),
        spec: Some(VmSpec {
            cpu: Some(plasmavmc_api::proto::CpuSpec {
                vcpus: 2,
                cores_per_socket: 1,
                sockets: 1,
                cpu_model: String::new(),
            }),
            memory: Some(plasmavmc_api::proto::MemorySpec {
                size_mib: 1024,  // 1 GB
                hugepages: false,
            }),
            disks: vec![],
            network: vec![],
            boot: None,
            security: None,
        }),
        hypervisor: ProtoHypervisorType::Kvm as i32,
        metadata: Default::default(),
        labels: Default::default(),
    }).await.unwrap().into_inner();

    assert!(!vm.id.is_empty());

    // 6. Verify balance was deducted after commit
    // Expected deduction: vcpus * 10 + memory_gb * 5 = 2 * 10 + 1 * 5 = 25
    let wallet_after = credit_client.get_wallet("proj2").await.unwrap();
    assert!(wallet_after.balance < 1000, "Balance should be reduced after VM creation");

    // 7. Cleanup: Delete VM
    let _ = vm_client.delete_vm(DeleteVmRequest {
        org_id: "org2".into(),
        project_id: "proj2".into(),
        vm_id: vm.id,
        force: true,
    }).await;

    // Clean up
    std::env::remove_var("CREDITSERVICE_ENDPOINT");
}

/// Test admission control without QEMU - uses mock/dry-run approach
/// This test validates the client integration code compiles and wires correctly
#[tokio::test]
async fn creditservice_client_integration_smoke() {
    // 1. Start CreditService
    let credit_addr = "127.0.0.1:50094";
    let storage: Arc<dyn CreditStorage> = InMemoryStorage::new();
    let credit_svc = CreditServiceImpl::new(storage.clone());

    let server_handle = tokio::spawn(async move {
        Server::builder()
            .add_service(CreditServiceServer::new(credit_svc))
            .serve(credit_addr.parse().unwrap())
            .await
            .unwrap();
    });
    tokio::time::sleep(std::time::Duration::from_millis(100)).await;

    // 2. Test CreditService client directly
    let mut client = CreditServiceClient::connect(format!("http://{}", credit_addr)).await.unwrap();

    // Create wallet
    let wallet = client.create_wallet("test-proj", "test-org", 500).await.unwrap();
    assert_eq!(wallet.project_id, "test-proj");
    assert_eq!(wallet.balance, 500);

    // Check quota (should pass)
    let check = client.check_quota(
        "test-proj",
        creditservice_client::ResourceType::VmInstance,
        1,
        100
    ).await.unwrap();
    assert!(check.allowed);

    // Reserve credits
    let reservation = client.reserve_credits(
        "test-proj",
        100,
        "Test VM creation",
        "VmInstance",
        300
    ).await.unwrap();
    assert!(!reservation.id.is_empty());

    // Commit reservation
    let commit = client.commit_reservation(&reservation.id, 100, "vm-123").await.unwrap();
    assert!(commit.transaction.is_some(), "Commit should create a transaction");

    // Verify balance reduced
    let wallet_after = client.get_wallet("test-proj").await.unwrap();
    assert_eq!(wallet_after.balance, 400); // 500 - 100

    // 3. Test reservation release (rollback)
    let reservation2 = client.reserve_credits(
        "test-proj",
        50,
        "Test VM creation 2",
        "VmInstance",
        300
    ).await.unwrap();

    // Release (rollback)
    let released = client.release_reservation(&reservation2.id, "cancelled").await.unwrap();
    assert!(released);

    // Balance should be unchanged after release
    let wallet_final = client.get_wallet("test-proj").await.unwrap();
    assert_eq!(wallet_final.balance, 400); // Still 400

    // Cleanup
    server_handle.abort();
}