millipede-storage-memory 0.1.0

In-memory StorageClient implementation for the Millipede web crawler.
Documentation
//! Request queue behavior tests.

use millipede_core::{
    request::Request,
    storage::{AddOptions, ReclaimOptions, RequestQueue, RequestSource, StorageError},
};
use millipede_storage_memory::{MemoryQueuePolicy, MemoryRequestQueue};
use std::time::Duration;

fn req(url: &str) -> Request {
    Request::get(url).build().unwrap()
}

#[tokio::test]
async fn dedup_tracks_pending_and_handled() {
    let queue = MemoryRequestQueue::new("dedup");
    let request = req("https://example.com/item");
    assert!(
        !queue
            .add(request.clone(), AddOptions::default())
            .await
            .unwrap()
            .was_already_present
    );
    let duplicate = queue
        .add(request.clone(), AddOptions::default())
        .await
        .unwrap();
    assert!(duplicate.was_already_present);
    assert!(!duplicate.was_already_handled);
    assert_eq!(queue.pending_count().await.unwrap(), 1);
    queue
        .mark_handled(queue.fetch_next().await.unwrap().unwrap())
        .await
        .unwrap();
    let duplicate = queue.add(request, AddOptions::default()).await.unwrap();
    assert!(duplicate.was_already_present && duplicate.was_already_handled);
    assert_eq!(queue.pending_count().await.unwrap(), 0);
}

#[tokio::test]
async fn forefront_precedes_fifo() {
    let queue = MemoryRequestQueue::new("front");
    for url in ["https://example.com/1", "https://example.com/2"] {
        let _ = queue.add(req(url), AddOptions::default()).await.unwrap();
    }
    let mut forefront = AddOptions::default();
    forefront.forefront = true;
    let _ = queue
        .add(req("https://example.com/3"), forefront)
        .await
        .unwrap();
    for path in ["/3", "/1", "/2"] {
        let lease = queue.fetch_next().await.unwrap().unwrap();
        assert_eq!(lease.request.url.path(), path);
        queue.mark_handled(lease).await.unwrap();
    }
}

#[tokio::test]
async fn lease_lifecycle_and_renewal() {
    let queue = MemoryRequestQueue::new("lease");
    let _ = queue
        .add(req("https://example.com/1"), AddOptions::default())
        .await
        .unwrap();
    let lease = queue.fetch_next().await.unwrap().unwrap();
    let lease_id = lease.lease_id.clone();
    assert!(queue.is_empty().await.unwrap());
    assert!(!queue.is_finished().await.unwrap());
    queue
        .renew(&lease_id, Duration::from_secs(30))
        .await
        .unwrap();
    queue.mark_handled(lease).await.unwrap();
    assert_eq!(queue.handled_count().await.unwrap(), 1);
    assert!(queue.is_finished().await.unwrap());
    assert!(matches!(
        queue.renew(&lease_id, Duration::from_secs(1)).await,
        Err(StorageError::LeaseNotFound { .. })
    ));
}

#[tokio::test]
async fn reclaim_controls_retry_and_position() {
    let queue = MemoryRequestQueue::new("reclaim");
    let first = req("https://example.com/first");
    let key = first.unique_key.clone();
    let _ = queue.add(first, AddOptions::default()).await.unwrap();
    let lease = queue.fetch_next().await.unwrap().unwrap();
    queue
        .reclaim(lease, ReclaimOptions::default())
        .await
        .unwrap();
    let lease = queue.fetch_next().await.unwrap().unwrap();
    assert_eq!(lease.request.unique_key, key);
    assert_eq!(lease.request.retry_count, 1);
    let mut no_increment = ReclaimOptions::default();
    no_increment.increment_retry = false;
    queue.reclaim(lease, no_increment).await.unwrap();
    let _ = queue
        .add(req("https://example.com/second"), AddOptions::default())
        .await
        .unwrap();
    let lease = queue.fetch_next().await.unwrap().unwrap();
    assert_eq!(lease.request.retry_count, 1);
    let mut forefront = ReclaimOptions::default();
    forefront.forefront = true;
    forefront.increment_retry = false;
    queue.reclaim(lease, forefront).await.unwrap();
    assert_eq!(
        queue
            .fetch_next()
            .await
            .unwrap()
            .unwrap()
            .request
            .unique_key,
        key
    );
}

#[tokio::test]
async fn abandon_does_not_increment_retry() {
    let queue = MemoryRequestQueue::new("abandon");
    let _ = queue
        .add(req("https://example.com/1"), AddOptions::default())
        .await
        .unwrap();
    queue
        .abandon(queue.fetch_next().await.unwrap().unwrap())
        .await
        .unwrap();
    assert_eq!(
        queue
            .fetch_next()
            .await
            .unwrap()
            .unwrap()
            .request
            .retry_count,
        0
    );
}

#[tokio::test]
async fn reclaim_persists_lease_request_mutations() {
    let queue = MemoryRequestQueue::new("reclaim-mutations");
    let _ = queue
        .add(req("https://example.com/1"), AddOptions::default())
        .await
        .unwrap();
    let mut lease = queue.fetch_next().await.unwrap().unwrap();
    lease.request.session_rotation_count = 3;
    lease
        .request
        .error_messages
        .push("temporary failure".into());
    queue
        .reclaim(lease, ReclaimOptions::default())
        .await
        .unwrap();

    let lease = queue.fetch_next().await.unwrap().unwrap();
    assert_eq!(lease.request.session_rotation_count, 3);
    assert_eq!(lease.request.error_messages, ["temporary failure"]);
    assert_eq!(lease.request.retry_count, 1);
}

#[tokio::test]
async fn abandon_persists_lease_request_mutations() {
    let queue = MemoryRequestQueue::new("abandon-mutations");
    let _ = queue
        .add(req("https://example.com/1"), AddOptions::default())
        .await
        .unwrap();
    let mut lease = queue.fetch_next().await.unwrap().unwrap();
    lease.request.session_rotation_count = 3;
    lease
        .request
        .error_messages
        .push("temporary failure".into());
    queue.abandon(lease).await.unwrap();

    let lease = queue.fetch_next().await.unwrap().unwrap();
    assert_eq!(lease.request.session_rotation_count, 3);
    assert_eq!(lease.request.error_messages, ["temporary failure"]);
    assert_eq!(lease.request.retry_count, 0);
}

#[tokio::test]
async fn add_batch_is_inline_and_flags_duplicate() {
    let queue = MemoryRequestQueue::new("batch");
    let duplicate = req("https://example.com/1");
    let handle = queue
        .add_batch(
            vec![
                duplicate.clone().into(),
                req("https://example.com/2").into(),
                RequestSource::Request(duplicate),
            ],
            AddOptions::default(),
        )
        .await
        .unwrap();
    assert_eq!(handle.added.len(), 3);
    assert!(handle.added[2].was_already_present);
    assert_eq!(handle.wait().await.unwrap().processed.len(), 3);
}

#[tokio::test]
async fn domain_round_robin_alternates() {
    let queue = MemoryRequestQueue::with_policy("fair", MemoryQueuePolicy::DomainRoundRobin);
    for url in ["https://a.com/1", "https://a.com/2", "https://b.com/1"] {
        let _ = queue.add(req(url), AddOptions::default()).await.unwrap();
    }
    for host in ["a.com", "b.com", "a.com"] {
        let lease = queue.fetch_next().await.unwrap().unwrap();
        assert_eq!(lease.request.url.host_str().unwrap(), host);
        queue.mark_handled(lease).await.unwrap();
    }
}