OpenShell/crates/openshell-server/src/test_support.rs

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Test fixtures for exercising gateway integration points.

use crate::ServerState;
use crate::auth::identity::{Identity, IdentityProvider};
use crate::auth::principal::{Principal, UserPrincipal};
use crate::compute::{new_test_runtime, new_test_runtime_for_driver};
use crate::persistence::Store;
use crate::sandbox_index::SandboxIndex;
use crate::sandbox_watch::SandboxWatchBus;
use crate::supervisor_session::SupervisorSessionRegistry;
use crate::tracing_bus::TracingLogBus;
use futures::{Stream, stream};
use openshell_core::Config;
#[cfg(unix)]
use openshell_core::proto::compute::v1::compute_driver_server::ComputeDriverServer;
use openshell_core::proto::compute::v1::{
    CreateSandboxRequest, CreateSandboxResponse, DeleteSandboxRequest, DeleteSandboxResponse,
    DeleteWorkspaceRequest, DeleteWorkspaceResponse, DriverSandbox, EnsureWorkspaceRequest,
    EnsureWorkspaceResponse, GetCapabilitiesRequest, GetCapabilitiesResponse, GetSandboxRequest,
    GetSandboxResponse, ListSandboxesRequest, ListSandboxesResponse, StartSandboxRequest,
    StartSandboxResponse, StopSandboxRequest, StopSandboxResponse, SuspendSandboxRequest,
    SuspendSandboxResponse, ValidateSandboxCreateRequest, ValidateSandboxCreateResponse,
    WatchSandboxesEvent, WatchSandboxesRequest, compute_driver_server::ComputeDriver,
};
use std::collections::HashMap;
#[cfg(unix)]
use std::io;
#[cfg(unix)]
use std::path::{Path, PathBuf};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
#[cfg(unix)]
use std::task::{Context, Poll};
#[cfg(unix)]
use tokio::net::{UnixListener, UnixStream};
#[cfg(unix)]
use tokio::task::JoinHandle;
use tonic::{Request, Response, Status};

type WatchStream = Pin<Box<dyn Stream<Item = Result<WatchSandboxesEvent, Status>> + Send>>;

/// Build a real in-memory gateway service backed by the requested test compute driver.
///
/// This fixture is intentionally narrow: integration tests can exercise the public
/// gRPC service without exposing the gateway's internal state construction details.
pub async fn gateway_service_with_driver(driver_name: &str) -> crate::OpenShellService {
    let store = Arc::new(
        Store::connect("sqlite::memory:?cache=shared")
            .await
            .expect("in-memory gateway store should open"),
    );
    crate::ensure_default_workspace(&store)
        .await
        .expect("default workspace should be created");
    let compute = if driver_name == "test" {
        new_test_runtime(store.clone()).await
    } else {
        new_test_runtime_for_driver(store.clone(), driver_name).await
    };
    let state = Arc::new(ServerState::new(
        Config::new(None)
            .with_database_url("sqlite::memory:?cache=shared")
            .with_credential_drivers(["test-static"]),
        store,
        compute,
        SandboxIndex::new(),
        SandboxWatchBus::new(),
        TracingLogBus::new(),
        Arc::new(SupervisorSessionRegistry::new()),
        None,
    ));
    crate::OpenShellService::new(state)
}

/// Tonic interceptor that authenticates integration-test requests as the dev user.
pub fn authenticate_as_dev_user(mut request: Request<()>) -> Result<Request<()>, Status> {
    request
        .extensions_mut()
        .insert(Principal::User(UserPrincipal {
            identity: Identity {
                subject: "dev-user".to_string(),
                display_name: None,
                roles: vec!["openshell-admin".to_string(), "openshell-user".to_string()],
                scopes: vec![],
                provider: IdentityProvider::Oidc,
            },
        }));
    Ok(request)
}

#[derive(Debug, Clone, PartialEq)]
pub enum FakeComputeDriverCall {
    GetCapabilities,
    ValidateSandboxCreate {
        sandbox: Option<DriverSandbox>,
    },
    GetSandbox {
        sandbox_id: String,
        sandbox_name: String,
    },
    ListSandboxes,
    CreateSandbox {
        sandbox: Option<DriverSandbox>,
    },
    UpdateSandboxConfiguration {
        request: openshell_core::proto::compute::v1::UpdateSandboxConfigurationRequest,
    },
    UpdateSandboxImage {
        request: openshell_core::proto::compute::v1::SandboxUpdatePreconditions,
    },
    StopSandbox {
        sandbox_id: String,
        sandbox_name: String,
    },
    StartSandbox {
        sandbox_id: String,
        sandbox_name: String,
        isolation_backend: Option<openshell_core::proto::compute::v1::IsolationBackendRegistration>,
        /// Fresh supervisor-facing credentials. These remain opaque in test
        /// diagnostics just as they do in production tracing.
        launch_authentication: Vec<u8>,
        generation_id: String,
        /// The optional durable provisioning snapshot supplied by newer
        /// gateways. Keeping it in the fake makes compatibility and
        /// reconciliation tests able to verify the start contract.
        sandbox: Option<DriverSandbox>,
    },
    DeleteSandbox {
        sandbox_id: String,
        sandbox_name: String,
    },
    WatchSandboxes,
}

#[derive(Debug, Clone)]
pub struct FakeComputeDriver {
    state: Arc<Mutex<FakeComputeDriverState>>,
}

#[derive(Debug)]
struct FakeComputeDriverState {
    capabilities: GetCapabilitiesResponse,
    sandboxes: HashMap<String, DriverSandbox>,
    calls: Vec<FakeComputeDriverCall>,
    traceparents: Vec<String>,
    image_update_error: Option<tonic::Code>,
    image_validation_error: Option<tonic::Code>,
    image_update_pause: Option<(Arc<tokio::sync::Notify>, Arc<tokio::sync::Semaphore>)>,
}

impl Default for FakeComputeDriver {
    fn default() -> Self {
        Self::new()
    }
}

impl FakeComputeDriver {
    /// Override the handshake response to exercise gateway capability admission.
    #[must_use]
    pub fn with_capabilities(self, capabilities: GetCapabilitiesResponse) -> Self {
        self.with_state(|state| state.capabilities = capabilities);
        self
    }

    #[must_use]
    pub fn new() -> Self {
        Self {
            state: Arc::new(Mutex::new(FakeComputeDriverState {
                capabilities: GetCapabilitiesResponse {
                    resource_admission_policy:
                        openshell_core::resource_admission::DriverAdmissionConfig::default()
                            .acknowledgement(),
                    driver_name: "fake-compute-driver".to_string(),
                    driver_version: "test".to_string(),
                    default_image: "openshell/sandbox:test".to_string(),
                    gateway_manages_lifecycle: false,
                    supports_sandbox_authentication: false,
                    driver_reports_runtime_readiness: false,
                    resource_capabilities: None,
                    rootfs_tar_staging_dir: String::new(),
                    rootfs_tar_max_bytes: 0,
                    extension: Some(openshell_core::extension_protocol::extension_metadata(
                        openshell_core::extension_protocol::ExtensionFamily::Compute,
                        "openshell/fake-compute-driver",
                        "test",
                        [],
                    )),
                },
                sandboxes: HashMap::new(),
                calls: Vec::new(),
                traceparents: Vec::new(),
                image_update_error: None,
                image_validation_error: None,
                image_update_pause: None,
            })),
        }
    }

    #[must_use]
    pub fn with_driver_name(self, driver_name: impl Into<String>) -> Self {
        self.with_state(|state| state.capabilities.driver_name = driver_name.into());
        self
    }

    #[must_use]
    pub fn with_driver_version(self, driver_version: impl Into<String>) -> Self {
        self.with_state(|state| state.capabilities.driver_version = driver_version.into());
        self
    }

    #[must_use]
    pub fn with_default_image(self, default_image: impl Into<String>) -> Self {
        self.with_state(|state| state.capabilities.default_image = default_image.into());
        self
    }

    #[must_use]
    pub fn with_gateway_manages_lifecycle(self) -> Self {
        self.with_state(|state| state.capabilities.gateway_manages_lifecycle = true);
        self
    }

    #[must_use]
    pub fn with_isolation_backend_registration(self) -> Self {
        self.with_state(|state| {
            state
                .capabilities
                .extension
                .as_mut()
                .unwrap()
                .supported_capabilities
                .push(openshell_core::extension_protocol::COMPUTE_ISOLATION_BACKEND.into());
        });
        self
    }

    #[must_use]
    pub fn with_sandbox_image_updates(self) -> Self {
        self.with_state(|state| {
            state
                .capabilities
                .extension
                .as_mut()
                .unwrap()
                .supported_capabilities
                .push(
                    openshell_core::extension_protocol::COMPUTE_UPDATE_SANDBOX_CONFIGURATION.into(),
                );
        });
        self
    }

    #[must_use]
    pub fn calls(&self) -> Vec<FakeComputeDriverCall> {
        self.with_state(|state| state.calls.clone())
    }

    #[must_use]
    pub fn traceparents(&self) -> Vec<String> {
        self.with_state(|state| state.traceparents.clone())
    }

    pub fn clear_calls(&self) {
        self.with_state(|state| state.calls.clear());
    }

    pub fn set_admission_acknowledgement(&self, acknowledgement: String) {
        self.with_state(|state| state.capabilities.resource_admission_policy = acknowledgement);
    }

    #[cfg(unix)]
    pub fn serve_uds(
        &self,
        socket_path: impl AsRef<Path>,
    ) -> io::Result<FakeComputeDriverServerHandle> {
        let socket_path = socket_path.as_ref().to_path_buf();
        let listener = UnixListener::bind(&socket_path)?;
        let driver = self.clone();
        let task = tokio::spawn(async move {
            tonic::transport::Server::builder()
                .add_service(ComputeDriverServer::new(driver.clone()))
                .add_service(openshell_core::proto::compute::v1::compute_driver_configuration_updates_server::ComputeDriverConfigurationUpdatesServer::new(driver.clone()))
                .serve_with_incoming(UnixIncoming { listener })
                .await
        });
        Ok(FakeComputeDriverServerHandle { socket_path, task })
    }

    pub fn pause_image_updates(&self) -> (Arc<tokio::sync::Notify>, Arc<tokio::sync::Semaphore>) {
        let entered = Arc::new(tokio::sync::Notify::new());
        let release = Arc::new(tokio::sync::Semaphore::new(0));
        self.with_state(|state| {
            state.image_update_pause = Some((entered.clone(), release.clone()));
        });
        (entered, release)
    }

    pub fn set_image_update_error(&self, code: Option<tonic::Code>) {
        self.with_state(|state| state.image_update_error = code);
    }

    pub fn set_image_validation_error(&self, code: Option<tonic::Code>) {
        self.with_state(|state| state.image_validation_error = code);
    }

    fn with_state<R>(&self, f: impl FnOnce(&mut FakeComputeDriverState) -> R) -> R {
        let mut state = self
            .state
            .lock()
            .expect("fake compute driver state poisoned");
        f(&mut state)
    }

    fn record_traceparent(&self, metadata: &tonic::metadata::MetadataMap) {
        let traceparent = metadata
            .get("traceparent")
            .and_then(|value| value.to_str().ok())
            .map(str::to_string);
        self.with_state(|state| state.traceparents.extend(traceparent));
    }
}

#[cfg(unix)]
#[derive(Debug)]
pub struct FakeComputeDriverServerHandle {
    socket_path: PathBuf,
    task: JoinHandle<Result<(), tonic::transport::Error>>,
}

#[cfg(unix)]
impl Drop for FakeComputeDriverServerHandle {
    fn drop(&mut self) {
        self.task.abort();
        let _ = std::fs::remove_file(&self.socket_path);
    }
}

#[cfg(unix)]
struct UnixIncoming {
    listener: UnixListener,
}

#[cfg(unix)]
impl Stream for UnixIncoming {
    type Item = io::Result<UnixStream>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        match self.get_mut().listener.poll_accept(cx) {
            Poll::Ready(Ok((stream, _addr))) => Poll::Ready(Some(Ok(stream))),
            Poll::Ready(Err(err)) => Poll::Ready(Some(Err(err))),
            Poll::Pending => Poll::Pending,
        }
    }
}

#[tonic::async_trait]
impl openshell_core::compute::ComputeDriver for FakeComputeDriver {
    async fn update_sandbox_configuration(
        &self,
        request: Request<openshell_core::proto::compute::v1::UpdateSandboxConfigurationRequest>,
    ) -> Result<
        Response<openshell_core::proto::compute::v1::UpdateSandboxConfigurationResponse>,
        Status,
    > {
        let pause = self.with_state(|state| state.image_update_pause.clone());
        if let Some((entered, release)) = pause {
            entered.notify_one();
            release.acquire().await.unwrap().forget();
        }
        self.record_traceparent(request.metadata());
        self.with_state(|state| {
            let request = request.into_inner();
            let validate_only = request.preconditions.as_ref().unwrap().validate_only;
            state.calls.push(FakeComputeDriverCall::UpdateSandboxImage {
                request: request.preconditions.clone().unwrap(),
            });
            state
                .calls
                .push(FakeComputeDriverCall::UpdateSandboxConfiguration { request });
            if let Some(code) = state.image_validation_error.filter(|_| validate_only) {
                return Err(Status::new(code, "injected configuration validation error"));
            }
            if let Some(code) = state.image_update_error.filter(|_| !validate_only) {
                return Err(Status::new(code, "injected configuration update error"));
            }
            Ok(Response::new(
                openshell_core::proto::compute::v1::UpdateSandboxConfigurationResponse {},
            ))
        })
    }
}

#[tonic::async_trait]
impl openshell_core::proto::compute::v1::compute_driver_configuration_updates_server::ComputeDriverConfigurationUpdates for FakeComputeDriver {
    async fn update_sandbox_configuration(&self, request: Request<openshell_core::proto::compute::v1::UpdateSandboxConfigurationRequest>) -> Result<Response<openshell_core::proto::compute::v1::UpdateSandboxConfigurationResponse>, Status> {
        openshell_core::compute::ComputeDriver::update_sandbox_configuration(self, request).await
    }
}

#[tonic::async_trait]
impl ComputeDriver for FakeComputeDriver {
    async fn authenticate_sandbox(
        &self,
        _request: Request<openshell_core::proto::compute::v1::AuthenticateSandboxRequest>,
    ) -> Result<Response<openshell_core::proto::compute::v1::AuthenticateSandboxResponse>, Status>
    {
        Err(Status::unimplemented(
            "fake driver does not authenticate sandbox credentials",
        ))
    }

    type WatchSandboxesStream = WatchStream;

    async fn get_capabilities(
        &self,
        request: Request<GetCapabilitiesRequest>,
    ) -> Result<Response<GetCapabilitiesResponse>, Status> {
        self.record_traceparent(request.metadata());
        let response = self.with_state(|state| {
            state.calls.push(FakeComputeDriverCall::GetCapabilities);
            state.capabilities.clone()
        });
        Ok(Response::new(response))
    }

    async fn validate_sandbox_create(
        &self,
        request: Request<ValidateSandboxCreateRequest>,
    ) -> Result<Response<ValidateSandboxCreateResponse>, Status> {
        self.record_traceparent(request.metadata());
        let request = request.into_inner();
        self.with_state(|state| {
            state
                .calls
                .push(FakeComputeDriverCall::ValidateSandboxCreate {
                    sandbox: request.sandbox,
                });
        });
        Ok(Response::new(ValidateSandboxCreateResponse {}))
    }

    async fn get_sandbox(
        &self,
        request: Request<GetSandboxRequest>,
    ) -> Result<Response<GetSandboxResponse>, Status> {
        self.record_traceparent(request.metadata());
        let request = request.into_inner();
        let sandbox = self.with_state(|state| {
            state.calls.push(FakeComputeDriverCall::GetSandbox {
                sandbox_id: request.sandbox_id.clone(),
                sandbox_name: request.name.clone(),
            });
            state
                .sandboxes
                .values()
                .find(|sandbox| {
                    (!request.sandbox_id.is_empty() && sandbox.id == request.sandbox_id)
                        || (!request.name.is_empty() && sandbox.name == request.name)
                })
                .cloned()
        });
        let sandbox = sandbox.ok_or_else(|| Status::not_found("sandbox not found"))?;
        Ok(Response::new(GetSandboxResponse {
            sandbox: Some(sandbox),
        }))
    }

    async fn list_sandboxes(
        &self,
        request: Request<ListSandboxesRequest>,
    ) -> Result<Response<ListSandboxesResponse>, Status> {
        self.record_traceparent(request.metadata());
        let sandboxes = self.with_state(|state| {
            state.calls.push(FakeComputeDriverCall::ListSandboxes);
            state.sandboxes.values().cloned().collect()
        });
        Ok(Response::new(ListSandboxesResponse { sandboxes }))
    }

    async fn create_sandbox(
        &self,
        request: Request<CreateSandboxRequest>,
    ) -> Result<Response<CreateSandboxResponse>, Status> {
        self.record_traceparent(request.metadata());
        let request = request.into_inner();
        let sandbox = request.sandbox;
        self.with_state(|state| {
            if let Some(sandbox) = sandbox.as_ref() {
                state.sandboxes.insert(sandbox.id.clone(), sandbox.clone());
            }
            state
                .calls
                .push(FakeComputeDriverCall::CreateSandbox { sandbox });
        });
        Ok(Response::new(CreateSandboxResponse::default()))
    }

    async fn suspend_sandbox(
        &self,
        _request: Request<SuspendSandboxRequest>,
    ) -> Result<Response<SuspendSandboxResponse>, Status> {
        Err(Status::unimplemented(
            "the test compute driver does not support sandbox suspend",
        ))
    }

    async fn stop_sandbox(
        &self,
        request: Request<StopSandboxRequest>,
    ) -> Result<Response<StopSandboxResponse>, Status> {
        self.record_traceparent(request.metadata());
        let request = request.into_inner();
        self.with_state(|state| {
            state.calls.push(FakeComputeDriverCall::StopSandbox {
                sandbox_id: request.sandbox_id,
                sandbox_name: request.name,
            });
        });
        Ok(Response::new(StopSandboxResponse {}))
    }

    async fn start_sandbox(
        &self,
        request: Request<StartSandboxRequest>,
    ) -> Result<Response<StartSandboxResponse>, Status> {
        self.record_traceparent(request.metadata());
        let request = request.into_inner();
        self.with_state(|state| {
            state.calls.push(FakeComputeDriverCall::StartSandbox {
                sandbox_id: request.sandbox_id,
                sandbox_name: request.name,
                isolation_backend: request.isolation_backend,
                launch_authentication: request.launch_authentication,
                generation_id: request.generation_id,
                sandbox: request.sandbox,
            });
        });
        Ok(Response::new(StartSandboxResponse::default()))
    }

    async fn delete_sandbox(
        &self,
        request: Request<DeleteSandboxRequest>,
    ) -> Result<Response<DeleteSandboxResponse>, Status> {
        self.record_traceparent(request.metadata());
        let request = request.into_inner();
        let deleted = self.with_state(|state| {
            state.calls.push(FakeComputeDriverCall::DeleteSandbox {
                sandbox_id: request.sandbox_id.clone(),
                sandbox_name: request.name.clone(),
            });
            if request.sandbox_id.is_empty() {
                let Some(id) = state
                    .sandboxes
                    .iter()
                    .find(|(_, sandbox)| sandbox.name == request.name)
                    .map(|(id, _)| id.clone())
                else {
                    return false;
                };
                state.sandboxes.remove(&id).is_some()
            } else {
                state.sandboxes.remove(&request.sandbox_id).is_some()
            }
        });
        Ok(Response::new(DeleteSandboxResponse { deleted }))
    }

    async fn watch_sandboxes(
        &self,
        request: Request<WatchSandboxesRequest>,
    ) -> Result<Response<Self::WatchSandboxesStream>, Status> {
        self.record_traceparent(request.metadata());
        self.with_state(|state| state.calls.push(FakeComputeDriverCall::WatchSandboxes));
        Ok(Response::new(Box::pin(stream::empty())))
    }

    async fn ensure_workspace(
        &self,
        _request: Request<EnsureWorkspaceRequest>,
    ) -> Result<Response<EnsureWorkspaceResponse>, Status> {
        Ok(Response::new(EnsureWorkspaceResponse {}))
    }

    async fn delete_workspace(
        &self,
        _request: Request<DeleteWorkspaceRequest>,
    ) -> Result<Response<DeleteWorkspaceResponse>, Status> {
        Ok(Response::new(DeleteWorkspaceResponse {}))
    }
}