notescraft/crates/engine/tests/sync_chunks.rs

//! Chunked uploads: the client against a server that keeps chunks in memory.
//!
//! The fake server holds the chunks it was sent in a map and can be told to
//! cut the connection, to lose a chunk or to change the file under the
//! client's feet.

use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::{Arc, Mutex, MutexGuard};

use notescraft_engine::{Changes, Limits, Progress, PutReply, Remote, SyncError, SyncOptions, SyncResult, SyncState, sync_once_with, upload_chunked};
use notescraft_sync::{Ledger, Outcome, Write, hash_hex};
use notescraft_vault::Vault;

/// The size of one chunk in the tests.
const CHUNK: usize = 1024;
/// The size above which a file uploads in chunks.
const THRESHOLD: u64 = 2 * 1024;

/// The server the client talks to.
struct Server {
    inner: Mutex<Inner>,
}

/// What the server keeps.
struct Inner {
    ledger: Ledger,
    /// The chunks the server holds, by their hash.
    chunks: BTreeMap<String, Vec<u8>>,
    /// The files the server holds, by their path.
    files: BTreeMap<String, Vec<u8>>,
    /// The hashes of every chunk that was sent.
    chunk_calls: Vec<String>,
    /// The chunk list of every chunked write.
    chunked_calls: Vec<Vec<String>>,
    /// How many chunk sends there were.
    calls: usize,
    /// From this send on, the connection is cut.
    cut_from: Option<usize>,
    /// The file that changes while the client is uploading it.
    corrupt: Option<(PathBuf, Vec<u8>)>,
}

impl Inner {
    fn new() -> Self {
        Self {
            ledger: Ledger::new(),
            chunks: BTreeMap::new(),
            files: BTreeMap::new(),
            chunk_calls: Vec::new(),
            chunked_calls: Vec::new(),
            calls: 0,
            cut_from: None,
            corrupt: None,
        }
    }
}

impl Server {
    fn new() -> Arc<Server> {
        Arc::new(Server { inner: Mutex::new(Inner::new()) })
    }

    fn lock(&self) -> MutexGuard<'_, Inner> {
        self.inner.lock().unwrap_or_else(std::sync::PoisonError::into_inner)
    }

    /// The hashes of the chunks that were sent.
    fn chunk_calls(&self) -> Vec<String> {
        self.lock().chunk_calls.clone()
    }

    /// How many chunked writes the server saw.
    fn chunked_calls(&self) -> usize {
        self.lock().chunked_calls.len()
    }

    /// The file the server holds.
    fn file(&self, path: &str) -> Option<Vec<u8>> {
        self.lock().files.get(path).cloned()
    }

    /// Cuts the connection from the given chunk send on.
    fn cut_from(&self, call: usize) {
        self.lock().cut_from = Some(call);
    }

    /// Lets the connection through again.
    fn mend(&self) {
        self.lock().cut_from = None;
    }

    /// Forgets the chunks that were sent, so a test can look at one pass only.
    fn forget_calls(&self) {
        self.lock().chunk_calls.clear();
    }

    /// Changes the file while the client uploads it.
    fn corrupt(&self, path: PathBuf, bytes: Vec<u8>) {
        self.lock().corrupt = Some((path, bytes));
    }
}

/// The client the driver talks to.
struct FakeRemote {
    server: Arc<Server>,
}

impl Remote for FakeRemote {
    fn changes(&self, since: u64, limit: usize, _wait: u64) -> SyncResult<Changes> {
        let inner = self.server.lock();
        Ok(inner.ledger.changes_since(since, limit))
    }

    fn get(&self, path: &str) -> SyncResult<Vec<u8>> {
        let inner = self.server.lock();
        inner.files.get(path).cloned().ok_or_else(|| SyncError::Status { status: 404, body: path.to_owned() })
    }

    fn put(&self, path: &str, bytes: &[u8], base_seq: u64, mtime_ms: u64) -> SyncResult<PutReply> {
        let mut inner = self.server.lock();
        let outcome = inner
            .ledger
            .apply(Write { path: path.to_owned(), hash: Some(hash_hex(bytes)), size: u64::try_from(bytes.len()).expect("size"), mtime_ms, base_seq })
            .map_err(|error| SyncError::Status { status: 400, body: error.to_string() })?;
        match outcome {
            Outcome::Applied(record) => {
                inner.files.insert(path.to_owned(), bytes.to_vec());
                Ok(PutReply::Applied(record))
            }
            Outcome::Unchanged(record) => Ok(PutReply::Unchanged(record)),
            Outcome::Conflict(record) => Ok(PutReply::Conflict(record)),
        }
    }

    fn delete(&self, path: &str, base_seq: u64) -> SyncResult<PutReply> {
        let mut inner = self.server.lock();
        let outcome = inner
            .ledger
            .apply(Write { path: path.to_owned(), hash: None, size: 0, mtime_ms: 0, base_seq })
            .map_err(|error| SyncError::Status { status: 400, body: error.to_string() })?;
        match outcome {
            Outcome::Applied(record) => {
                inner.files.remove(path);
                Ok(PutReply::Applied(record))
            }
            Outcome::Unchanged(record) => Ok(PutReply::Unchanged(record)),
            Outcome::Conflict(record) => Ok(PutReply::Conflict(record)),
        }
    }

    fn missing_chunks(&self, hashes: &[String]) -> SyncResult<Vec<String>> {
        let mut inner = self.server.lock();
        if let Some((path, bytes)) = inner.corrupt.take() {
            std::fs::write(&path, &bytes).map_err(|error| SyncError::Io(error.to_string()))?;
        }
        Ok(hashes.iter().filter(|hash| !inner.chunks.contains_key(*hash)).cloned().collect())
    }

    fn put_chunk(&self, hash: &str, bytes: &[u8]) -> SyncResult<()> {
        let mut inner = self.server.lock();
        inner.calls = inner.calls.saturating_add(1);
        inner.chunk_calls.push(hash.to_owned());
        if inner.cut_from.is_some_and(|from| inner.calls >= from) {
            return Err(SyncError::Http("the connection was cut".to_owned()));
        }
        inner.chunks.insert(hash.to_owned(), bytes.to_vec());
        Ok(())
    }

    fn put_chunked(&self, path: &str, chunk_hashes: &[String], content_hash: &str, base_seq: u64, mtime_ms: u64) -> SyncResult<PutReply> {
        let mut inner = self.server.lock();
        inner.chunked_calls.push(chunk_hashes.to_vec());
        let mut bytes = Vec::new();
        for hash in chunk_hashes {
            let Some(chunk) = inner.chunks.get(hash) else { return Err(SyncError::MissingChunks(vec![hash.clone()])) };
            bytes.extend_from_slice(chunk);
        }
        let outcome = inner
            .ledger
            .apply(Write { path: path.to_owned(), hash: Some(content_hash.to_owned()), size: u64::try_from(bytes.len()).expect("size"), mtime_ms, base_seq })
            .map_err(|error| SyncError::Status { status: 400, body: error.to_string() })?;
        match outcome {
            Outcome::Applied(record) => {
                inner.files.insert(path.to_owned(), bytes);
                Ok(PutReply::Applied(record))
            }
            Outcome::Unchanged(record) => Ok(PutReply::Unchanged(record)),
            Outcome::Conflict(record) => Ok(PutReply::Conflict(record)),
        }
    }
}

/// A vault in a fresh folder.
fn temp_vault(name: &str) -> Vault {
    let root = std::env::temp_dir().join(format!("nc-engine-chunks-{name}-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&root);
    std::fs::create_dir_all(&root).expect("temp folder");
    Vault::open(&root).expect("vault")
}

/// The limits the tests use: small chunks, so a small file has many of them.
fn limits() -> Limits {
    Limits { threshold: THRESHOLD, chunk: CHUNK }
}

/// Ten kilobytes whose chunks differ from one another.
fn ten_chunks() -> Vec<u8> {
    (0..10u8).flat_map(|index| vec![index; CHUNK]).collect()
}

/// The hash of the chunk number `index` of `bytes`.
fn chunk_hash(bytes: &[u8], index: usize) -> String {
    let start = index * CHUNK;
    let end = start.saturating_add(CHUNK);
    hash_hex(bytes.get(start..end).expect("chunk"))
}

#[test]
fn a_big_file_travels_in_chunks_and_is_written_once() {
    let server = Server::new();
    let vault = temp_vault("big");
    let bytes = ten_chunks();
    vault.write("big.bin", &bytes).expect("write");
    let remote = FakeRemote { server: server.clone() };

    let reply = upload_chunked(&remote, &vault, "big.bin", 0, 0, limits(), &mut |_: &Progress| {}).expect("upload");

    assert!(matches!(reply, PutReply::Applied(_)), "{reply:?}");
    assert_eq!(server.chunk_calls().len(), 10, "one send per chunk");
    assert_eq!(server.chunked_calls(), 1, "the file is written once");
    assert_eq!(server.file("big.bin"), Some(bytes), "the server holds the whole file");
}

#[test]
fn a_cut_connection_resumes_with_the_chunks_that_never_arrived() {
    let server = Server::new();
    let vault = temp_vault("resume");
    let bytes = ten_chunks();
    vault.write("big.bin", &bytes).expect("write");
    let remote = FakeRemote { server: server.clone() };

    // The connection dies on the sixth chunk; the chunk is tried again and fails.
    server.cut_from(6);
    let error = upload_chunked(&remote, &vault, "big.bin", 0, 0, limits(), &mut |_: &Progress| {}).expect_err("cut");
    assert!(matches!(error, SyncError::Http(_)), "{error:?}");
    assert_eq!(server.chunked_calls(), 0, "the file was never written");

    // The next pass sends only the chunks that never arrived.
    server.mend();
    server.forget_calls();
    upload_chunked(&remote, &vault, "big.bin", 0, 0, limits(), &mut |_: &Progress| {}).expect("resumed");

    let sent: Vec<String> = (5..10).map(|index| chunk_hash(&bytes, index)).collect();
    assert_eq!(server.chunk_calls(), sent, "only the missing chunks are sent again");
    assert_eq!(server.file("big.bin"), Some(bytes), "the server holds the whole file");
}

#[test]
fn the_same_chunk_of_a_file_is_sent_once() {
    let server = Server::new();
    let vault = temp_vault("same");
    let bytes = vec![7u8; CHUNK * 4];
    vault.write("same.bin", &bytes).expect("write");
    let remote = FakeRemote { server: server.clone() };

    upload_chunked(&remote, &vault, "same.bin", 0, 0, limits(), &mut |_: &Progress| {}).expect("upload");

    assert_eq!(server.chunk_calls().len(), 1, "the four equal chunks travel once");
    assert_eq!(server.file("same.bin"), Some(bytes), "the file still holds the four blocks");
}

#[test]
fn the_progress_grows_and_ends_at_the_size_of_the_file() {
    let server = Server::new();
    let vault = temp_vault("progress");
    let bytes = ten_chunks();
    vault.write("big.bin", &bytes).expect("write");
    let remote = FakeRemote { server: server.clone() };

    let mut seen: Vec<Progress> = Vec::new();
    upload_chunked(&remote, &vault, "big.bin", 0, 0, limits(), &mut |progress: &Progress| seen.push(progress.clone())).expect("upload");

    assert_eq!(seen.len(), 10, "one report per chunk");
    assert!(seen.windows(2).all(|pair| pair[0].done_bytes <= pair[1].done_bytes), "the progress never goes back");
    let last = seen.last().expect("a report");
    assert_eq!(last.done_bytes, last.total_bytes);
    assert_eq!(last.total_bytes, u64::try_from(bytes.len()).expect("size"));
    assert_eq!(last.path, "big.bin");
}

#[test]
fn a_small_file_never_touches_the_chunk_methods() {
    let server = Server::new();
    let vault = temp_vault("small");
    vault.write("small.md", b"# one\n").expect("write");
    let remote = FakeRemote { server: server.clone() };
    let mut state = SyncState::default();

    let report = sync_once_with(&vault, &remote, &mut state, "phone", "2026-10-20", SyncOptions { limits: limits(), progress: Box::new(|_: &Progress| {}) })
        .expect("sync");

    assert_eq!(report.uploaded, vec!["small.md".to_owned()]);
    assert_eq!(server.chunk_calls(), Vec::<String>::new(), "a small file is not chunked");
    assert_eq!(server.chunked_calls(), 0);
    assert_eq!(server.file("small.md"), Some(b"# one\n".to_vec()));
}

#[test]
fn a_file_that_changes_while_it_is_uploaded_is_refused() {
    let server = Server::new();
    let vault = temp_vault("changed");
    let bytes = ten_chunks();
    vault.write("big.bin", &bytes).expect("write");
    let remote = FakeRemote { server: server.clone() };
    // The file changes between the hash pass and the upload pass.
    server.corrupt(vault.root().join("big.bin"), vec![1u8; CHUNK * 10]);

    let error = upload_chunked(&remote, &vault, "big.bin", 0, 0, limits(), &mut |_: &Progress| {}).expect_err("changed");

    assert!(matches!(error, SyncError::BadRemote(_)), "{error:?}");
    assert_eq!(server.chunked_calls(), 0, "no mixed file reaches the server");
    assert_eq!(server.file("big.bin"), None);
}

#[test]
fn an_empty_file_goes_as_a_plain_write() {
    let server = Server::new();
    let vault = temp_vault("empty");
    vault.write("empty.bin", b"").expect("write");
    let remote = FakeRemote { server: server.clone() };

    let reply = upload_chunked(&remote, &vault, "empty.bin", 0, 0, limits(), &mut |_: &Progress| {}).expect("upload");

    assert!(matches!(reply, PutReply::Applied(_)), "{reply:?}");
    assert_eq!(server.chunk_calls(), Vec::<String>::new(), "an empty file has no chunk");
    assert_eq!(server.chunked_calls(), 0);
    assert_eq!(server.file("empty.bin"), Some(Vec::new()));
}

#[test]
fn the_threshold_decides_what_is_chunked() {
    let server = Server::new();
    let vault = temp_vault("threshold");
    let at_threshold: Vec<u8> = (0..THRESHOLD).map(|index| u8::try_from(index % 251).expect("byte")).collect();
    let over_threshold: Vec<u8> = (0..=THRESHOLD).map(|index| u8::try_from(index % 251).expect("byte")).collect();
    vault.write("at.bin", &at_threshold).expect("write");
    vault.write("over.bin", &over_threshold).expect("write");
    let remote = FakeRemote { server: server.clone() };
    let mut state = SyncState::default();

    let report = sync_once_with(&vault, &remote, &mut state, "phone", "2026-10-20", SyncOptions { limits: limits(), progress: Box::new(|_: &Progress| {}) })
        .expect("sync");

    assert_eq!(report.uploaded.len(), 2);
    assert_eq!(server.chunked_calls(), 1, "only the file over the threshold is chunked");
    assert_eq!(server.chunk_calls().len(), 3, "the file over the threshold has three chunks");
    assert_eq!(server.file("at.bin"), Some(at_threshold));
    assert_eq!(server.file("over.bin"), Some(over_threshold));
}