sq/crates/sq-sim/tests/scenarios/single_node.rs

use std::path::PathBuf;
use std::sync::Arc;

use sq_models::WalConfig;
use sq_sim::fs::InMemoryFileSystem;
use sq_sim::SimClock;
use sq_storage::engine::StorageEngine;

fn test_engine() -> (
    StorageEngine<InMemoryFileSystem, SimClock>,
    Arc<InMemoryFileSystem>,
    Arc<SimClock>,
) {
    let fs = Arc::new(InMemoryFileSystem::new());
    let clock = Arc::new(SimClock::new());
    let config = WalConfig {
        max_segment_bytes: 1024 * 1024,
        max_segment_age_secs: 3600,
        data_dir: PathBuf::from("/data"),
        ..Default::default()
    };
    let engine = StorageEngine::new(fs.clone(), clock.clone(), config).unwrap();
    (engine, fs, clock)
}

/// S01: Single node, single producer, single consumer - baseline correctness
#[test]
fn s01_single_producer_consumer() {
    let (engine, _fs, _clock) = test_engine();

    // Produce 1000 messages.
    for i in 0..1000u64 {
        let offset = engine
            .append("orders", 0, Some(format!("key-{i}").as_bytes()), format!("value-{i}").as_bytes(), &[], i)
            .unwrap();
        assert_eq!(offset, i, "offset must match sequence");
    }

    // Consume all messages.
    let messages = engine.read("orders", 0, 0, 2000).unwrap();

    // Invariant 1: No message loss.
    assert_eq!(messages.len(), 1000);

    // Invariant 2: Offsets strictly monotonic, no gaps.
    for (i, msg) in messages.iter().enumerate() {
        assert_eq!(msg.offset, i as u64, "offset gap detected at index {i}");
    }

    // Invariant: Content integrity.
    for msg in &messages {
        let expected_key = format!("key-{}", msg.offset);
        let expected_value = format!("value-{}", msg.offset);
        assert_eq!(msg.key.as_ref().unwrap(), expected_key.as_bytes());
        assert_eq!(msg.value, expected_value.as_bytes());
    }
}

/// S02: Single node, concurrent producers to different topics - offset ordering
#[test]
fn s02_multi_topic_producers() {
    let (engine, _fs, _clock) = test_engine();

    let topics = ["events", "orders", "logs"];

    // Write 100 messages to each topic.
    for topic in &topics {
        for i in 0..100u64 {
            let offset = engine.append(topic, 0, None, b"data", &[], i).unwrap();
            assert_eq!(offset, i);
        }
    }

    // Verify each topic has its own offset space.
    for topic in &topics {
        let messages = engine.read(topic, 0, 0, 200).unwrap();
        assert_eq!(messages.len(), 100, "topic {topic} should have 100 messages");

        // Offsets are monotonic per topic.
        for (i, msg) in messages.iter().enumerate() {
            assert_eq!(msg.offset, i as u64);
        }
    }

    // Cross-topic isolation: reading one topic doesn't return messages from another.
    let events = engine.read("events", 0, 0, 200).unwrap();
    for msg in &events {
        assert_eq!(msg.topic.as_str(), "events");
    }
}

/// S03: Single node, disk full during write - graceful error handling
#[test]
fn s03_disk_full() {
    let fs = Arc::new(InMemoryFileSystem::new());
    let clock = Arc::new(SimClock::new());
    let config = WalConfig {
        max_segment_bytes: 1024 * 1024,
        max_segment_age_secs: 3600,
        data_dir: PathBuf::from("/data"),
        ..Default::default()
    };
    let engine = StorageEngine::new(fs.clone(), clock, config).unwrap();

    // Write some messages successfully.
    for i in 0..10 {
        engine.append("t", 0, None, b"data", &[], i).unwrap();
    }

    // Simulate disk full.
    fs.simulate_disk_full();

    // Next write should fail.
    let result = engine.append("t", 0, None, b"data", &[], 0);
    assert!(result.is_err(), "write should fail when disk is full");

    // Clear fault - subsequent writes should work.
    fs.clear_faults();
    let _offset = engine.append("t", 0, None, b"after-recovery", &[], 0).unwrap();

    // Verify earlier messages are still readable.
    let messages = engine.read("t", 0, 0, 100).unwrap();
    assert!(messages.len() >= 10, "original messages should survive disk full");
}

/// S04: Single node, crash and restart - WAL recovery
#[test]
fn s04_crash_recovery() {
    let fs = Arc::new(InMemoryFileSystem::new());
    let clock = Arc::new(SimClock::new());
    let config = WalConfig {
        max_segment_bytes: 1024 * 1024,
        max_segment_age_secs: 3600,
        data_dir: PathBuf::from("/data"),
        ..Default::default()
    };

    // Phase 1: Write messages and "crash" (drop engine).
    {
        let engine = StorageEngine::new(fs.clone(), clock.clone(), config.clone()).unwrap();
        for i in 0..500u64 {
            engine
                .append("orders", 0, None, format!("msg-{i}").as_bytes(), &[], i)
                .unwrap();
        }
        // Engine dropped here - simulates crash.
    }

    // Phase 2: "Restart" - create new engine and recover.
    {
        let engine = StorageEngine::new(fs.clone(), clock.clone(), config.clone()).unwrap();
        engine.recover().unwrap();

        // Invariant 1: All acked messages survive recovery.
        let messages = engine.read("orders", 0, 0, 1000).unwrap();
        assert_eq!(messages.len(), 500, "all messages must survive crash");

        // Invariant 2: Offsets are intact.
        for (i, msg) in messages.iter().enumerate() {
            assert_eq!(msg.offset, i as u64);
            assert_eq!(msg.value, format!("msg-{i}").as_bytes());
        }

        // Can continue writing after recovery.
        let offset = engine.append("orders", 0, None, b"post-crash", &[], 0).unwrap();
        assert_eq!(offset, 500);
    }
}

/// S09: Consumer group offset preservation across restarts
#[test]
fn s09_consumer_group_offset_persistence() {
    let fs = Arc::new(InMemoryFileSystem::new());
    let clock = Arc::new(SimClock::new());
    let config = WalConfig {
        max_segment_bytes: 1024 * 1024,
        max_segment_age_secs: 3600,
        data_dir: PathBuf::from("/data"),
        ..Default::default()
    };

    // Write messages and commit an offset.
    {
        let engine = StorageEngine::new(fs.clone(), clock.clone(), config.clone()).unwrap();
        for i in 0..100 {
            engine.append("t", 0, None, b"data", &[], i).unwrap();
        }
        engine.commit_offset("group-1", "t", 0, 50).unwrap();
    }

    // Restart and verify committed offset survives.
    {
        let engine = StorageEngine::new(fs.clone(), clock.clone(), config.clone()).unwrap();
        engine.recover().unwrap();

        // Invariant 4: Consumer group offsets never regress.
        let committed = engine.get_committed_offset("group-1", "t", 0);
        assert_eq!(committed, Some(50));

        // Can resume consuming from committed offset.
        let messages = engine.read("t", 0, 51, 100).unwrap();
        assert_eq!(messages.len(), 49); // offsets 51-99
    }
}

/// S10: High throughput burst - no message loss
#[test]
fn s10_high_throughput() {
    let (engine, _fs, _clock) = test_engine();

    let msg_count = 10_000u64;

    // Burst write.
    for i in 0..msg_count {
        engine
            .append("burst", 0, None, format!("msg-{i}").as_bytes(), &[], i)
            .unwrap();
    }

    // Verify no loss.
    let messages = engine.read("burst", 0, 0, (msg_count + 1) as usize).unwrap();
    assert_eq!(messages.len(), msg_count as usize);

    // Verify ordering.
    for (i, msg) in messages.iter().enumerate() {
        assert_eq!(msg.offset, i as u64);
    }
}

/// S06: Segment rotation and recovery - multiple segments survive crash
#[test]
fn s06_segment_rotation_recovery() {
    let fs = Arc::new(InMemoryFileSystem::new());
    let clock = Arc::new(SimClock::new());
    let config = WalConfig {
        max_segment_bytes: 512, // Very small segments to force rotation.
        max_segment_age_secs: 3600,
        data_dir: PathBuf::from("/data"),
        ..Default::default()
    };

    // Write enough messages to cause multiple segment rotations.
    {
        let engine = StorageEngine::new(fs.clone(), clock.clone(), config.clone()).unwrap();
        for i in 0..200u64 {
            engine
                .append("t", 0, None, format!("msg-{i}").as_bytes(), &[], i)
                .unwrap();
        }
    }

    // Recover.
    {
        let engine = StorageEngine::new(fs.clone(), clock.clone(), config.clone()).unwrap();
        engine.recover().unwrap();

        let messages = engine.read("t", 0, 0, 300).unwrap();
        assert_eq!(messages.len(), 200, "all messages across segments must survive");

        for (i, msg) in messages.iter().enumerate() {
            assert_eq!(msg.offset, i as u64);
        }

        // Continue writing.
        let offset = engine.append("t", 0, None, b"new", &[], 0).unwrap();
        assert_eq!(offset, 200);
    }
}