//! 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, /// `ETag`, verbatim, quotes included. pub etag: String, /// `X-Linked-Etag`, present for LFS-backed files. pub linked_etag: Option, pub linked_size: Option, } 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, } /// The fake Hub's mutable state. #[derive(Default)] pub struct HubState { /// `repo_id` -> ref name -> commit. pub refs: HashMap>, /// `repo_id` -> commit -> revision. pub revisions: HashMap>, /// Repos that answer 403 with `X-Error-Code: GatedRepo`. pub gated: Vec, /// 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>, /// Every `GET` on the resolve path. This is the number single-flight is /// asserted against. pub file_gets: Arc, /// Every request of any kind, so a sealed test can prove zero contact. pub requests: Arc, shutdown: Option>, } 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::("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>, file_gets: Arc, requests: Arc, } /// 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, Path(rest): Path) -> 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::>(), })) .into_response() } /// The fake Hub's resolve path. async fn resolve( State(hub): State, method: Method, Path(rest): Path, ) -> 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 { 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, 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 = Arc::new(ObjectStoreAdapter::local(bucket).expect("local store")); let upstream: Arc = 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 }