memona-plugin-aindrive/src/session_tests.rs

//! 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};
use crate::sdk::{EntryKind, ErrorCode, Guest, GuestReader, Locator, Setting};
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 && !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<u8>, Option<String>), 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<u8> = (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());
}