rustingface/crates/rustingface-api/tests/support/mod.rs

//! A fake Hub, and a rustingface wired to it.
//!
//! The tests that matter are about behaviour against a real client protocol,
//! so the fake upstream speaks the same headers the Hub does — `X-Repo-Commit`,
//! `X-Linked-Etag`, `X-Linked-Size`, `X-Error-Code` — and counts what it was
//! asked for, which is how single-flight is proved rather than assumed.

#![allow(dead_code)]

use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use axum::Router;
use axum::body::Body;
use axum::extract::{Path, State};
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use http::{HeaderMap, Method, StatusCode};
use rustingface_api::auth::Tokens;
use rustingface_api::{AppState, router};
use rustingface_core::Registry;
use rustingface_core::ports::{Store, Upstream};
use rustingface_data::{HubClient, ObjectStoreAdapter, SealedUpstream};
use rustingface_entities::config::{AuthMode, Config};
use tokio::sync::RwLock;

/// One file the fake Hub serves.
#[derive(Clone)]
pub struct HubFile {
    pub body: Vec<u8>,
    /// `ETag`, verbatim, quotes included.
    pub etag: String,
    /// `X-Linked-Etag`, present for LFS-backed files.
    pub linked_etag: Option<String>,
    pub linked_size: Option<u64>,
}

impl HubFile {
    /// A small text file, the shape `config.json` takes upstream: a git blob
    /// etag and no LFS headers.
    pub fn plain(body: &str) -> Self {
        Self {
            body: body.as_bytes().to_vec(),
            etag: format!("\"{:040x}\"", body.len()),
            linked_etag: None,
            linked_size: None,
        }
    }

    /// An LFS-backed file, whose linked etag is the sha256 of its bytes.
    pub fn lfs(body: &[u8]) -> Self {
        use sha2::{Digest, Sha256};
        let oid = hex::encode(Sha256::digest(body));
        Self {
            body: body.to_vec(),
            etag: format!("\"{oid}\""),
            linked_etag: Some(format!("\"{oid}\"")),
            linked_size: Some(body.len() as u64),
        }
    }

    /// An LFS-backed file whose advertised oid does not match its bytes.
    pub fn corrupt(body: &[u8]) -> Self {
        let mut file = Self::lfs(body);
        file.linked_etag = Some(format!("\"{}\"", "b".repeat(64)));
        file
    }
}

/// One revision of one repo.
#[derive(Clone, Default)]
pub struct HubRevision {
    pub files: HashMap<String, HubFile>,
}

/// The fake Hub's mutable state.
#[derive(Default)]
pub struct HubState {
    /// `repo_id` -> ref name -> commit.
    pub refs: HashMap<String, HashMap<String, String>>,
    /// `repo_id` -> commit -> revision.
    pub revisions: HashMap<String, HashMap<String, HubRevision>>,
    /// Repos that answer 403 with `X-Error-Code: GatedRepo`.
    pub gated: Vec<String>,
    /// Delay inserted between body chunks, to keep a transfer observable.
    pub chunk_delay: Duration,
    /// Bytes per body chunk.
    pub chunk_size: usize,
}

/// A running fake Hub.
pub struct FakeHub {
    pub base_url: String,
    pub state: Arc<RwLock<HubState>>,
    /// Every `GET` on the resolve path. This is the number single-flight is
    /// asserted against.
    pub file_gets: Arc<AtomicUsize>,
    /// Every request of any kind, so a sealed test can prove zero contact.
    pub requests: Arc<AtomicUsize>,
    shutdown: Option<tokio::sync::oneshot::Sender<()>>,
}

impl FakeHub {
    /// Start a fake Hub on an ephemeral port.
    pub async fn start() -> Self {
        let state = Arc::new(RwLock::new(HubState {
            chunk_size: 64 * 1024,
            ..Default::default()
        }));
        let file_gets = Arc::new(AtomicUsize::new(0));
        let requests = Arc::new(AtomicUsize::new(0));

        let hub = HubHandle {
            state: Arc::clone(&state),
            file_gets: Arc::clone(&file_gets),
            requests: Arc::clone(&requests),
        };

        let app = Router::new()
            .route("/api/models/{*rest}", get(api))
            .route("/api/datasets/{*rest}", get(api))
            .route("/{*rest}", get(resolve).head(resolve))
            .with_state(hub);

        let listener = tokio::net::TcpListener::bind::<SocketAddr>("127.0.0.1:0".parse().unwrap())
            .await
            .expect("binding the fake hub");
        let addr = listener.local_addr().unwrap();
        let (tx, rx) = tokio::sync::oneshot::channel();
        tokio::spawn(async move {
            let _ = axum::serve(listener, app)
                .with_graceful_shutdown(async {
                    let _ = rx.await;
                })
                .await;
        });

        Self {
            base_url: format!("http://{addr}"),
            state,
            file_gets,
            requests,
            shutdown: Some(tx),
        }
    }

    /// Publish a revision and point a ref at it.
    pub async fn publish(
        &self,
        repo: &str,
        ref_name: &str,
        commit: &str,
        files: &[(&str, HubFile)],
    ) {
        let mut state = self.state.write().await;
        state
            .refs
            .entry(repo.to_owned())
            .or_default()
            .insert(ref_name.to_owned(), commit.to_owned());
        let revision = HubRevision {
            files: files
                .iter()
                .map(|(path, file)| ((*path).to_owned(), file.clone()))
                .collect(),
        };
        state
            .revisions
            .entry(repo.to_owned())
            .or_default()
            .insert(commit.to_owned(), revision);
    }

    /// Move a ref to a different commit, as an upstream force-push would.
    pub async fn move_ref(&self, repo: &str, ref_name: &str, commit: &str) {
        self.state
            .write()
            .await
            .refs
            .entry(repo.to_owned())
            .or_default()
            .insert(ref_name.to_owned(), commit.to_owned());
    }

    /// Make a repo answer 403 `GatedRepo`.
    pub async fn gate(&self, repo: &str) {
        self.state.write().await.gated.push(repo.to_owned());
    }

    /// Slow the body down so a transfer stays observable mid-flight.
    pub async fn throttle(&self, chunk_size: usize, delay: Duration) {
        let mut state = self.state.write().await;
        state.chunk_size = chunk_size;
        state.chunk_delay = delay;
    }

    /// `GET`s on the resolve path so far.
    pub fn file_gets(&self) -> usize {
        self.file_gets.load(Ordering::SeqCst)
    }

    /// Requests of any kind so far.
    pub fn requests(&self) -> usize {
        self.requests.load(Ordering::SeqCst)
    }

    /// Forget everything counted so far.
    pub fn reset_counts(&self) {
        self.file_gets.store(0, Ordering::SeqCst);
        self.requests.store(0, Ordering::SeqCst);
    }

    /// Stop answering, as null-routing the Hub at the site router would.
    pub fn stop(&mut self) {
        if let Some(tx) = self.shutdown.take() {
            let _ = tx.send(());
        }
    }
}

#[derive(Clone)]
struct HubHandle {
    state: Arc<RwLock<HubState>>,
    file_gets: Arc<AtomicUsize>,
    requests: Arc<AtomicUsize>,
}

/// Reject a gated repo the way the Hub does.
fn gated() -> Response {
    (
        StatusCode::FORBIDDEN,
        [("x-error-code", "GatedRepo")],
        "access to this repo is gated",
    )
        .into_response()
}

/// Report a missing entity with the code the client branches on.
fn not_found(code: &'static str) -> Response {
    (StatusCode::NOT_FOUND, [("x-error-code", code)], code).into_response()
}

/// The fake Hub's `/api/{type}/…` surface.
async fn api(State(hub): State<HubHandle>, Path(rest): Path<String>) -> Response {
    hub.requests.fetch_add(1, Ordering::SeqCst);
    let state = hub.state.read().await;
    let segments: Vec<&str> = rest.split('/').collect();

    let marker = (1..=2)
        .filter(|i| *i < segments.len())
        .find(|i| matches!(segments[*i], "revision" | "tree"));
    let Some(i) = marker else {
        return not_found("RepoNotFound");
    };
    let repo = segments[..i].join("/");
    let revision = segments.get(i + 1).copied().unwrap_or("main");

    if state.gated.contains(&repo) {
        return gated();
    }
    let Some(commit) = resolve_commit(&state, &repo, revision) else {
        return not_found("RevisionNotFound");
    };
    let Some(rev) = state.revisions.get(&repo).and_then(|r| r.get(&commit)) else {
        return not_found("RevisionNotFound");
    };

    let mut paths: Vec<&String> = rev.files.keys().collect();
    paths.sort();

    if segments[i] == "tree" {
        let tree: Vec<_> = paths
            .iter()
            .map(|path| {
                let file = &rev.files[*path];
                serde_json::json!({
                    "type": "file",
                    "path": path,
                    "size": file.body.len(),
                    "oid": file.etag.trim_matches('"'),
                })
            })
            .collect();
        return axum::Json(tree).into_response();
    }

    axum::Json(serde_json::json!({
        "id": repo,
        "sha": commit,
        "lastModified": "2026-06-11T00:00:00.000Z",
        "private": false,
        "gated": false,
        "tags": ["text-generation"],
        "siblings": paths.iter().map(|p| serde_json::json!({"rfilename": p})).collect::<Vec<_>>(),
    }))
    .into_response()
}

/// The fake Hub's resolve path.
async fn resolve(
    State(hub): State<HubHandle>,
    method: Method,
    Path(rest): Path<String>,
) -> Response {
    hub.requests.fetch_add(1, Ordering::SeqCst);
    let state = hub.state.read().await;

    let rest = rest.strip_prefix("datasets/").unwrap_or(&rest);
    let Some((repo, tail)) = rest.split_once("/resolve/") else {
        return not_found("RepoNotFound");
    };
    if state.gated.contains(&repo.to_owned()) {
        return gated();
    }
    let Some((revision, path)) = tail.split_once('/') else {
        return not_found("EntryNotFound");
    };
    let revision = percent_decode(revision);

    let Some(commit) = resolve_commit(&state, repo, &revision) else {
        return not_found("RevisionNotFound");
    };
    let Some(file) = state
        .revisions
        .get(repo)
        .and_then(|r| r.get(&commit))
        .and_then(|r| r.files.get(path))
        .cloned()
    else {
        return not_found("EntryNotFound");
    };

    let mut headers = HeaderMap::new();
    headers.insert("x-repo-commit", commit.parse().unwrap());
    headers.insert("etag", file.etag.parse().unwrap());
    if let Some(linked) = &file.linked_etag {
        headers.insert("x-linked-etag", linked.parse().unwrap());
    }
    if let Some(size) = file.linked_size {
        headers.insert("x-linked-size", size.to_string().parse().unwrap());
    }
    headers.insert(
        "content-length",
        file.body.len().to_string().parse().unwrap(),
    );

    if method == Method::HEAD {
        return (StatusCode::OK, headers).into_response();
    }
    hub.file_gets.fetch_add(1, Ordering::SeqCst);

    let chunk_size = state.chunk_size.max(1);
    let delay = state.chunk_delay;
    let body = file.body.clone();
    drop(state);

    let stream = futures::stream::unfold(0usize, move |offset| {
        let body = body.clone();
        async move {
            if offset >= body.len() {
                return None;
            }
            if !delay.is_zero() {
                tokio::time::sleep(delay).await;
            }
            let end = (offset + chunk_size).min(body.len());
            let chunk = bytes::Bytes::copy_from_slice(&body[offset..end]);
            Some((Ok::<_, std::io::Error>(chunk), end))
        }
    });

    (StatusCode::OK, headers, Body::from_stream(stream)).into_response()
}

/// Resolve a ref or pass a commit through.
fn resolve_commit(state: &HubState, repo: &str, revision: &str) -> Option<String> {
    if let Some(commit) = state.refs.get(repo).and_then(|r| r.get(revision)) {
        return Some(commit.clone());
    }
    state
        .revisions
        .get(repo)
        .and_then(|r| r.contains_key(revision).then(|| revision.to_owned()))
}

/// Undo the percent-encoding the resolve URL applies to a revision.
fn percent_decode(value: &str) -> String {
    let mut out = String::with_capacity(value.len());
    let bytes = value.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i] == b'%' && i + 2 < bytes.len() {
            if let Ok(byte) = u8::from_str_radix(&value[i + 1..i + 3], 16) {
                out.push(byte as char);
                i += 3;
                continue;
            }
        }
        out.push(bytes[i] as char);
        i += 1;
    }
    out
}

/// A rustingface instance under test, over a temporary bucket directory.
pub struct Instance {
    pub app: axum::Router,
    pub registry: Registry,
    pub store: Arc<dyn Store>,
    pub config: Config,
    pub bucket: std::path::PathBuf,
}

/// Build an instance against `hub`, applying `tweak` to the config first.
pub fn instance(
    bucket: &std::path::Path,
    hub: Option<&FakeHub>,
    tweak: impl FnOnce(&mut Config),
) -> Instance {
    let mut config = Config::default();
    config.storage.local_path = Some(bucket.to_owned());
    config.storage.multipart_part_size = 5 * 1024 * 1024;
    config.server.client_stall_timeout = Duration::from_secs(2);
    match hub {
        Some(hub) => {
            config.upstream.enabled = true;
            config.upstream.endpoint = hub.base_url.clone();
        }
        None => config.upstream.enabled = false,
    }
    tweak(&mut config);
    config.validate().expect("test config must be valid");

    let store: Arc<dyn Store> = Arc::new(ObjectStoreAdapter::local(bucket).expect("local store"));
    let upstream: Arc<dyn Upstream> = if config.upstream.enabled {
        Arc::new(HubClient::new(&config.upstream, None).expect("hub client"))
    } else {
        Arc::new(SealedUpstream::new(config.upstream.endpoint.clone()))
    };

    let registry = Registry::new(Arc::clone(&store), upstream, config.clone()).expect("registry");
    let tokens = Tokens::new(config.auth.mode, None)
        .unwrap_or_else(|_| Tokens::new(AuthMode::None, None).expect("permissive tokens"));
    let app = router(AppState::new(registry.clone(), Arc::new(tokens), None));

    Instance {
        app,
        registry,
        store,
        config,
        bucket: bucket.to_owned(),
    }
}

/// A fresh directory to use as a bucket.
pub fn bucket_dir(name: &str) -> std::path::PathBuf {
    let dir = std::env::temp_dir().join(format!(
        "rustingface-test-{name}-{}-{}",
        std::process::id(),
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    ));
    std::fs::create_dir_all(&dir).expect("creating the test bucket");
    dir
}