//! The provider against an in-memory aindrive. Live QA separately runs the //! compiled component through Memona's host against the real service. use super::{Aindrive, SESSION, download, test_host, upload}; use crate::sdk::{ EntryKind, ErrorCode, Guest, GuestReader, GuestWriter, Locator, Setting, WriteCondition, }; use std::collections::BTreeMap; const ENDPOINT: &str = "https://aindrive.ainetwork.ai/mcp/d/KQHGpl8hFXq8"; fn config() -> Vec<(String, Setting)> { vec![ ( "url".to_owned(), Setting::Text("https://aindrive.ainetwork.ai/d/KQHGpl8hFXq8".to_owned()), ), ( "token".to_owned(), Setting::Secret(" aind_pat_secret\n".to_owned()), ), ] } fn at(path: &str) -> Locator { Locator::Path(path.to_owned()) } fn sample() -> test_host::Drive { let mut files = BTreeMap::new(); for (path, bytes, mtime) in [ ("PAGE.md", &b"# Studio"[..], "1791731214739.939"), ("notes/작가노트_2026.md", &b"notes"[..], "1791726579100"), ("notes/odd\\name.md", &b"x"[..], "1"), ("artworks.tsv", &b"title\tyear"[..], "1791726579103.648"), ] { files.insert(path.to_owned(), test_host::file(bytes, mtime)); } test_host::Drive { files, folders: vec!["outputs".to_owned()], ..test_host::Drive::default() } } #[tokio::test] async fn opening_checks_the_token_against_the_drive_endpoint() -> Result<(), ErrorCode> { let drive = test_host::install(sample()); let capabilities = Aindrive::open("aindrive".to_owned(), config())?; assert!(capabilities.write && capabilities.remove); assert!(capabilities.conditional_write); assert!(!capabilities.mkdir && !capabilities.rename && !capabilities.watch); let drive = drive.borrow(); let [seen] = drive.seen.as_slice() else { return Err(ErrorCode::Internal); }; assert_eq!(seen.url, ENDPOINT); assert_eq!(seen.method, "tools/list"); assert_eq!( seen.authorization.as_deref(), Some("Bearer aind_pat_secret") ); Ok(()) } #[tokio::test] async fn a_refused_token_or_a_missing_drive_leaves_no_connection() { for (status, expected) in [(401, ErrorCode::Denied), (404, ErrorCode::NotFound)] { let _drive = test_host::install(test_host::Drive { status: Some(status), ..sample() }); assert_eq!( Aindrive::open("aindrive".to_owned(), config()).err(), Some(expected) ); assert!(SESSION.with_borrow(Option::is_none)); } } #[tokio::test] async fn a_token_that_cannot_read_files_is_refused() { let _drive = test_host::install(test_host::Drive { tools: Some(vec!["search"]), ..sample() }); assert_eq!( Aindrive::open("aindrive".to_owned(), config()).err(), Some(ErrorCode::Denied) ); assert!(SESSION.with_borrow(Option::is_none)); } #[tokio::test] async fn a_config_without_a_token_or_drive_is_invalid() { let _drive = test_host::install(sample()); let only_url = config().into_iter().take(1).collect(); let only_token = config().into_iter().skip(1).collect(); for config in [only_url, only_token] { assert_eq!( Aindrive::open("aindrive".to_owned(), config).err(), Some(ErrorCode::InvalidInput) ); } } #[tokio::test] async fn listing_pages_through_a_folder_and_hides_unshowable_names() -> Result<(), ErrorCode> { let _drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; let first = Aindrive::list(at("/"), None, 2)?; let names: Vec<_> = first.entries.iter().map(|e| e.name.as_str()).collect(); assert_eq!(names, ["PAGE.md", "artworks.tsv"]); assert_eq!( first.entries[0].info.revision.as_deref(), Some("m1791731214739.939-s8") ); let rest = Aindrive::list(at("/"), first.cursor, 2)?; let names: Vec<_> = rest.entries.iter().map(|e| e.name.as_str()).collect(); assert_eq!(names, ["notes", "outputs"]); assert!(rest.cursor.is_none()); assert!(matches!(rest.entries[0].info.kind, EntryKind::Directory)); assert_eq!(rest.entries[0].info.size, 0); assert!(rest.entries[0].info.revision.is_none()); let notes = Aindrive::list(at("/notes/"), None, 10)?; let names: Vec<_> = notes.entries.iter().map(|e| e.name.as_str()).collect(); assert_eq!(names, ["작가노트_2026.md"]); Ok(()) } #[tokio::test] async fn listing_a_file_or_a_missing_folder_says_which() -> Result<(), ErrorCode> { let _drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; assert_eq!( Aindrive::list(at("/PAGE.md"), None, 10).err(), Some(ErrorCode::InvalidInput) ); assert_eq!( Aindrive::list(at("/gone"), None, 10).err(), Some(ErrorCode::NotFound) ); assert_eq!( Aindrive::list(at("/"), Some("x".to_owned()), 10).err(), Some(ErrorCode::InvalidInput) ); Ok(()) } #[tokio::test] async fn stat_reports_the_root_files_and_folders() -> Result<(), ErrorCode> { let _drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; let root = Aindrive::stat(at("/"))?; assert!(matches!(root.kind, EntryKind::Directory)); let page = Aindrive::stat(at("/PAGE.md"))?; assert!(matches!(page.kind, EntryKind::File)); assert_eq!(page.size, 8); assert_eq!(page.revision.as_deref(), Some("m1791731214739.939-s8")); assert!(matches!( Aindrive::stat(at("/outputs"))?.kind, EntryKind::Directory )); assert_eq!( Aindrive::stat(at("/missing.md")).err(), Some(ErrorCode::NotFound) ); Ok(()) } /// Reads all of a file from `offset`, as the host does, in pulls of /// `limit` bytes that may come back shorter. fn read_from( path: &str, offset: u64, limit: u32, ) -> Result<(Vec, Option), ErrorCode> { let reader = download(&at(path), offset)?; let mut read = Vec::new(); loop { let piece = reader.pull(limit)?; assert!(piece.len() <= limit as usize); if piece.is_empty() { return Ok((read, reader.revision())); } read.extend_from_slice(&piece); } } #[tokio::test] async fn reads_start_anywhere_and_never_exceed_the_limit() -> Result<(), ErrorCode> { // Longer than one of the test host's 4 KiB reply pieces, so groups of // base64 characters are split between pieces. let bytes: Vec = (0..=255).cycle().take(20_000).collect(); let mut drive = sample(); drive .files .insert("data.bin".to_owned(), test_host::file(&bytes, "5")); let _drive = test_host::install(drive); Aindrive::open("aindrive".to_owned(), config())?; for offset in [0_u64, 1, 2, 3, 3070, 3071, 10_000, 19_998, 19_999] { for limit in [1, 7, 4096, 1024 * 1024] { let (read, revision) = read_from("/data.bin", offset, limit)?; assert_eq!(read, bytes[offset as usize..], "{offset} by {limit}"); assert_eq!(revision.as_deref(), Some("m5-s20000")); } } assert!(read_from("/data.bin", 20_000, 10)?.0.is_empty()); assert!(read_from("/data.bin", 50_000, 10)?.0.is_empty()); Ok(()) } #[tokio::test] async fn files_shorter_than_a_group_and_empty_files_read_whole() -> Result<(), ErrorCode> { let mut drive = sample(); for (name, bytes) in [("one", &b"a"[..]), ("two", b"ab"), ("empty", b"")] { drive .files .insert(name.to_owned(), test_host::file(bytes, "1")); } let _drive = test_host::install(drive); Aindrive::open("aindrive".to_owned(), config())?; assert_eq!(read_from("/one", 0, 10)?.0, b"a"); assert_eq!(read_from("/two", 1, 10)?.0, b"b"); assert_eq!(read_from("/empty", 0, 10)?.0, b""); Ok(()) } #[tokio::test] async fn a_file_changed_between_stat_and_read_is_a_conflict() -> Result<(), ErrorCode> { let _drive = test_host::install(test_host::Drive { change_on_read: Some(test_host::file(b"# Studio, longer now", "2")), ..sample() }); Aindrive::open("aindrive".to_owned(), config())?; assert_eq!( download(&at("/PAGE.md"), 0).err(), Some(ErrorCode::Conflict) ); Ok(()) } #[tokio::test] async fn folders_missing_files_and_oversized_files_are_not_read() -> Result<(), ErrorCode> { let mut drive = sample(); drive.files.insert( "large.bin".to_owned(), test_host::file(&vec![0; 8 * 1024 * 1024 + 1], "9"), ); let drive = test_host::install(drive); Aindrive::open("aindrive".to_owned(), config())?; assert_eq!( download(&at("/notes"), 0).err(), Some(ErrorCode::InvalidInput) ); assert_eq!( download(&at("/missing.md"), 0).err(), Some(ErrorCode::NotFound) ); drive.borrow_mut().seen.clear(); assert_eq!( download(&at("/large.bin"), 0).err(), Some(ErrorCode::Unsupported) ); let tools: Vec<_> = drive .borrow() .seen .iter() .map(|seen| seen.tool.clone()) .collect(); assert_eq!(tools, [Some("stat".to_owned())]); Ok(()) } #[tokio::test] async fn nothing_is_requested_before_a_connection_opens() { let drive = test_host::install(sample()); Aindrive::close(); assert_eq!(Aindrive::stat(at("/")).err(), Some(ErrorCode::Internal)); assert!(drive.borrow().seen.is_empty()); } /// Saves `bytes` the way the host does: one push, then one finish. fn save( path: &str, condition: WriteCondition, bytes: &[u8], ) -> Result, ErrorCode> { let writer = upload(&at(path), condition, Some(bytes.len() as u64))?; writer.push(bytes.to_vec())?; let info = writer.finish()?.ok_or(ErrorCode::Internal)?; assert_eq!(info.size, bytes.len() as u64); Ok(info.revision) } #[tokio::test] async fn a_read_only_token_opens_a_read_only_folder() -> Result<(), ErrorCode> { let _drive = test_host::install(test_host::Drive { tools: Some(vec!["list_files", "read_file", "stat", "search"]), ..sample() }); let capabilities = Aindrive::open("aindrive".to_owned(), config())?; assert!(!capabilities.write && !capabilities.remove); assert!(!capabilities.conditional_write); Ok(()) } #[tokio::test] async fn a_save_on_the_revision_read_lands_and_reads_back() -> Result<(), ErrorCode> { let _drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; let (_, read) = read_from("/PAGE.md", 0, 4096)?; let read = read.ok_or(ErrorCode::Internal)?; let saved = save("/PAGE.md", WriteCondition::IfMatch(read), b"# Saved")?; let (bytes, revision) = read_from("/PAGE.md", 0, 4096)?; assert_eq!(bytes, b"# Saved"); assert_eq!(revision, saved); assert_eq!(Aindrive::stat(at("/PAGE.md"))?.revision, saved); Ok(()) } #[tokio::test] async fn a_save_on_a_stale_revision_is_a_conflict_and_changes_nothing() -> Result<(), ErrorCode> { let _drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; let stale = WriteCondition::IfMatch("m1-s8".to_owned()); assert_eq!(save("/PAGE.md", stale, b"lost"), Err(ErrorCode::Conflict)); let gone = WriteCondition::IfMatch("m1-s8".to_owned()); assert_eq!(save("/missing.md", gone, b"x"), Err(ErrorCode::Conflict)); assert_eq!(read_from("/PAGE.md", 0, 4096)?.0, b"# Studio"); Ok(()) } #[tokio::test] async fn create_only_refuses_anything_already_there() -> Result<(), ErrorCode> { let _drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; for taken in ["/PAGE.md", "/notes"] { assert_eq!( save(taken, WriteCondition::CreateOnly, b"new"), Err(ErrorCode::Conflict), "{taken}" ); } assert!( save("/outputs/new.md", WriteCondition::CreateOnly, b"new")?.is_some() ); assert_eq!(read_from("/outputs/new.md", 0, 4096)?.0, b"new"); Ok(()) } #[tokio::test] async fn names_that_need_escaping_and_every_byte_survive_a_save() -> Result<(), ErrorCode> { let _drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; let bytes: Vec = (0..=255).rev().cycle().take(70_001).collect(); let path = "/notes/\"quoted\" 노트 {x}.bin"; save(path, WriteCondition::Overwrite, &bytes)?; assert_eq!(read_from(path, 0, 1024 * 1024)?.0, bytes); save("/empty.md", WriteCondition::Overwrite, b"")?; assert_eq!(read_from("/empty.md", 0, 10)?.0, b""); Ok(()) } #[tokio::test] async fn a_save_larger_than_one_transfer_is_refused_before_it_is_sent() -> Result<(), ErrorCode> { let drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; let over = Some(1024 * 1024 + 1); assert_eq!( upload(&at("/big.bin"), WriteCondition::Overwrite, over).err(), Some(ErrorCode::Unsupported) ); let writer = upload(&at("/big.bin"), WriteCondition::Overwrite, None)?; writer.push(vec![1; 1024 * 1024])?; assert_eq!(writer.push(vec![1]), Err(ErrorCode::Unsupported)); let tools: Vec<_> = drive.borrow().seen.iter().map(|s| s.tool.clone()).collect(); assert_eq!(tools, [None]); Ok(()) } #[tokio::test] async fn an_abandoned_or_finished_save_sends_nothing_more() -> Result<(), ErrorCode> { let drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; let abandoned = upload(&at("/PAGE.md"), WriteCondition::Overwrite, None)?; abandoned.push(b"half".to_vec())?; drop(abandoned); let writer = upload(&at("/PAGE.md"), WriteCondition::Overwrite, None)?; writer.push(b"whole".to_vec())?; writer.finish()?; assert_eq!(writer.finish().err(), Some(ErrorCode::InvalidInput)); assert_eq!(writer.push(b"more".to_vec()), Err(ErrorCode::InvalidInput)); let writes = drive .borrow() .seen .iter() .filter(|seen| seen.tool.as_deref() == Some("write_file")) .count(); assert_eq!(writes, 1); assert_eq!(read_from("/PAGE.md", 0, 4096)?.0, b"whole"); Ok(()) } #[tokio::test] async fn a_save_whose_answer_is_lost_has_an_unknown_outcome() -> Result<(), ErrorCode> { let _drive = test_host::install(test_host::Drive { lose_reply: Some("write_file"), ..sample() }); Aindrive::open("aindrive".to_owned(), config())?; assert_eq!( save("/PAGE.md", WriteCondition::Overwrite, b"landed"), Err(ErrorCode::OutcomeUnknown) ); // It did land, which is why a retry must look first. assert_eq!(read_from("/PAGE.md", 0, 4096)?.0, b"landed"); Ok(()) } #[tokio::test] async fn a_device_that_goes_quiet_leaves_saves_unknown_and_reads_unavailable() -> Result<(), ErrorCode> { for tool in ["write_file", "stat"] { let _drive = test_host::install(test_host::Drive { fail: Some((tool, "[internal] agent timeout")), ..sample() }); Aindrive::open("aindrive".to_owned(), config())?; let expected = match tool { "write_file" => ErrorCode::OutcomeUnknown, _ => ErrorCode::Unavailable, }; assert_eq!( save("/PAGE.md", WriteCondition::Overwrite, b"x").err(), (tool == "write_file").then_some(expected) ); assert_eq!( Aindrive::stat(at("/PAGE.md")).err(), (tool == "stat").then_some(expected) ); } let _drive = test_host::install(test_host::Drive { fail: Some(("write_file", "[internal] agent offline")), ..sample() }); Aindrive::open("aindrive".to_owned(), config())?; assert_eq!( save("/PAGE.md", WriteCondition::Overwrite, b"x"), Err(ErrorCode::Unavailable) ); Ok(()) } #[tokio::test] async fn files_are_removed_and_folders_are_left_alone() -> Result<(), ErrorCode> { let drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; Aindrive::remove(at("/artworks.tsv"))?; assert_eq!( Aindrive::stat(at("/artworks.tsv")).err(), Some(ErrorCode::NotFound) ); for folder in ["/notes", "/outputs"] { assert_eq!( Aindrive::remove(at(folder)).err(), Some(ErrorCode::Unsupported), "{folder}" ); } assert!(drive.borrow().files.contains_key("notes/작가노트_2026.md")); assert_eq!( Aindrive::remove(at("/missing.md")).err(), Some(ErrorCode::NotFound) ); assert_eq!( Aindrive::remove(at("/")).err(), Some(ErrorCode::InvalidInput) ); Ok(()) } #[tokio::test] async fn only_a_server_error_leaves_a_save_unknown() -> Result<(), ErrorCode> { for (status, expected) in [ (502, ErrorCode::OutcomeUnknown), (429, ErrorCode::Unavailable), (403, ErrorCode::Denied), ] { let drive = test_host::install(sample()); Aindrive::open("aindrive".to_owned(), config())?; drive.borrow_mut().status = Some(status); assert_eq!( save("/PAGE.md", WriteCondition::Overwrite, b"x"), Err(expected), "{status}" ); assert_eq!( Aindrive::stat(at("/PAGE.md")).err(), Some(match status { 502 | 429 => ErrorCode::Unavailable, _ => expected, }) ); } Ok(()) }