remora/crates/remora-result/tests/substrate.rs

use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use remora_domain::OpaqueId;
use remora_durable::{PoolConfig, connect, migrate as durable_migrate};
use remora_result::{
    MANIFEST_SCHEMA_VERSION, PageSize, ResultError, ResultInventory, ResultMember,
    migrate as result_migrate,
};
use remora_stream::TargetBody;
use remora_target::{
    ContentSha256, CreateObjectRequest, ObjectIdentity, StoredObject, TargetStore,
    TargetStoreConfig, hash_bytes,
};
use serde_json::Value;
use sqlx::PgPool;

static NEXT_ID: AtomicU64 = AtomicU64::new(0);

fn unique_id(kind: &str) -> OpaqueId {
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("system time")
        .as_nanos();
    let sequence = NEXT_ID.fetch_add(1, Ordering::Relaxed);
    OpaqueId::parse(format!("{kind}-{nanos:x}-{sequence:x}")).expect("opaque id")
}

async fn test_pool() -> Option<PgPool> {
    let Some(url) = std::env::var("DATABASE_URL")
        .ok()
        .filter(|value| !value.is_empty())
    else {
        if std::env::var("REMORA_RESULT_REQUIRE").is_ok() {
            panic!("DATABASE_URL is required when REMORA_RESULT_REQUIRE is set");
        }
        return None;
    };
    let pool = connect(
        &url,
        PoolConfig {
            max_connections: 8,
            acquire_timeout: Duration::from_secs(10),
            idle_timeout: Some(Duration::from_secs(30)),
        },
    )
    .await
    .expect("connect postgres");
    durable_migrate(&pool).await.expect("durable migration");
    result_migrate(&pool).await.expect("result migration");
    Some(pool)
}

fn target_store() -> Option<TargetStore> {
    if std::env::var("REMORA_TARGET_ENDPOINT").is_err() {
        if std::env::var("REMORA_RESULT_REQUIRE").is_ok() {
            panic!("target fixture is required when REMORA_RESULT_REQUIRE is set");
        }
        return None;
    }
    TargetStore::connect(TargetStoreConfig::from_env().expect("target config"))
        .map(Some)
        .expect("connect target")
}

async fn stored_member(
    target: &TargetStore,
    attempt: &OpaqueId,
    path: &str,
    body: &[u8],
) -> ResultMember {
    let member_id = unique_id("member");
    let stored = target
        .put_create_only(
            &CreateObjectRequest::for_attempt_object(attempt, &member_id, None),
            body.to_vec(),
        )
        .await
        .expect("store member");
    ResultMember::new(
        member_id,
        path,
        stored.size,
        stored.identity,
        stored.content_sha256,
    )
    .expect("member")
}

#[tokio::test]
async fn single_member_manifest_resolves_pages_and_exact_access() {
    let (Some(pool), Some(target)) = (test_pool().await, target_store()) else {
        return;
    };
    let inventory = ResultInventory::new(pool);
    let attempt = unique_id("attempt");
    inventory.create(&attempt).await.expect("create inventory");
    let member = stored_member(&target, &attempt, "models/é.bin", b"member bytes").await;
    inventory
        .add_member(&attempt, &member)
        .await
        .expect("add member");

    let committed = inventory
        .create_manifest(&attempt, &unique_id("manifest"), &target)
        .await
        .expect("create manifest");
    assert_eq!(committed.schema_version, MANIFEST_SCHEMA_VERSION);
    assert_eq!(committed.member_count, 1);
    assert_eq!(committed.total_bytes, 12);
    inventory
        .validate_publication_input(&attempt, &committed)
        .await
        .expect("valid publication input");

    let manifest_bytes = target
        .read_exact(&committed.manifest.identity)
        .await
        .expect("read exact manifest");
    assert_eq!(
        hash_bytes(&manifest_bytes),
        committed.manifest.content_sha256
    );
    let manifest: Value = serde_json::from_slice(&manifest_bytes).expect("manifest JSON");
    assert_eq!(manifest["schema_version"], 1);
    assert_eq!(manifest["member_count"], 1);
    assert_eq!(manifest["members"][0]["path"], "models/é.bin");

    let page = inventory
        .page_members(&committed, None, PageSize::new(10).unwrap())
        .await
        .expect("page members");
    assert_eq!(page.items.len(), 1);
    assert_eq!(page.items[0], member);
    assert!(page.next_after.is_none());
    let found = inventory
        .lookup_member(&committed, member.id())
        .await
        .expect("lookup exact member");
    assert_eq!(found, member);
}

#[tokio::test]
async fn rejects_normalization_collision_and_closed_inventory_mutation() {
    let Some(pool) = test_pool().await else {
        return;
    };
    let inventory = ResultInventory::new(pool);
    let attempt = unique_id("attempt");
    inventory.create(&attempt).await.expect("create inventory");
    let first = ResultMember::new(
        unique_id("member"),
        "café/file",
        1,
        member_identity(&attempt, &unique_id("object")),
        hash_bytes(b"a"),
    )
    .unwrap();
    let collision = ResultMember::new(
        unique_id("member"),
        "cafe\u{301}/file",
        1,
        member_identity(&attempt, &unique_id("object")),
        hash_bytes(b"b"),
    )
    .unwrap();
    inventory.add_member(&attempt, &first).await.unwrap();
    assert!(matches!(
        inventory.add_member(&attempt, &collision).await,
        Err(ResultError::PathCollision)
    ));
    inventory.close(&attempt).await.unwrap();
    assert!(matches!(
        inventory.add_member(&attempt, &collision).await,
        Err(ResultError::InventoryNotOpen)
    ));
}

#[tokio::test]
async fn integrity_stale_attempt_and_descriptor_version_fail_closed() {
    let (Some(pool), Some(target)) = (test_pool().await, target_store()) else {
        return;
    };
    let inventory = ResultInventory::new(pool);
    let attempt = unique_id("attempt");
    inventory.create(&attempt).await.unwrap();
    let stored = stored_member(&target, &attempt, "payload.bin", b"payload").await;
    let invalid = ResultMember::new(
        stored.id().clone(),
        stored.path(),
        stored.size() + 1,
        stored.exact_identity(),
        stored.content_sha256().clone(),
    )
    .unwrap();
    inventory.add_member(&attempt, &invalid).await.unwrap();
    assert!(matches!(
        inventory
            .create_manifest(&attempt, &unique_id("manifest"), &target)
            .await,
        Err(ResultError::Integrity)
    ));

    let other_attempt = unique_id("attempt");
    let stale = ResultMember::new(
        unique_id("member"),
        "stale.bin",
        0,
        member_identity(&other_attempt, &unique_id("object")),
        hash_bytes(b""),
    )
    .unwrap();
    assert!(matches!(
        inventory.add_member(&attempt, &stale).await,
        Err(ResultError::StaleAttempt)
    ));

    let mut unsupported = placeholder_publication(&attempt);
    unsupported.schema_version = MANIFEST_SCHEMA_VERSION + 1;
    assert!(matches!(
        inventory.resolve_committed(&unsupported).await,
        Err(ResultError::UnsupportedManifestVersion(_))
    ));
}

#[tokio::test]
async fn pages_large_ordered_inventory_without_materializing_all_rows() {
    let (Some(pool), Some(target)) = (test_pool().await, target_store()) else {
        return;
    };
    let inventory = ResultInventory::new(pool);
    let attempt = unique_id("attempt");
    inventory.create(&attempt).await.unwrap();
    for index in (0..2_500_u32).rev() {
        let member_id = OpaqueId::parse(format!("member-{index:08}")).unwrap();
        let stored = target
            .put_create_only(
                &CreateObjectRequest::for_attempt_object(&attempt, &member_id, None),
                Vec::new(),
            )
            .await
            .unwrap();
        let member = ResultMember::new(
            member_id,
            &format!("inventory/{index:08}"),
            stored.size,
            stored.identity,
            stored.content_sha256,
        )
        .unwrap();
        inventory.add_member(&attempt, &member).await.unwrap();
    }
    let (committed, metrics) = inventory
        .create_manifest_with_metrics(&attempt, &unique_id("manifest"), &target)
        .await
        .unwrap();
    assert_eq!(metrics.members_streamed, 2_500);
    assert!(metrics.max_encoded_member_bytes < 1_024);
    assert_eq!(committed.member_count, 2_500);

    let mut after = None;
    let mut observed = 0_u32;
    loop {
        let page = inventory
            .page_members(&committed, after.as_deref(), PageSize::new(137).unwrap())
            .await
            .unwrap();
        for member in &page.items {
            assert_eq!(member.path(), format!("inventory/{observed:08}"));
            observed += 1;
        }
        let Some(next) = page.next_after else {
            break;
        };
        after = Some(next);
    }
    assert_eq!(observed, 2_500);
}

#[tokio::test]
async fn binary_path_order_is_identical_in_manifest_and_pages() {
    let (Some(pool), Some(target)) = (test_pool().await, target_store()) else {
        return;
    };
    let inventory = ResultInventory::new(pool);
    let attempt = unique_id("attempt");
    inventory.create(&attempt).await.unwrap();
    for path in ["한/file", "é/file", "a/file", "Z/file"] {
        let member = stored_member(&target, &attempt, path, b"").await;
        inventory.add_member(&attempt, &member).await.unwrap();
    }
    let committed = inventory
        .create_manifest(&attempt, &unique_id("manifest"), &target)
        .await
        .unwrap();
    let manifest: Value = serde_json::from_slice(
        &target
            .read_exact(&committed.manifest.identity)
            .await
            .unwrap(),
    )
    .unwrap();
    let manifest_paths = manifest["members"]
        .as_array()
        .unwrap()
        .iter()
        .map(|member| member["path"].as_str().unwrap())
        .collect::<Vec<_>>();
    assert_eq!(manifest_paths, ["Z/file", "a/file", "é/file", "한/file"]);
    let page = inventory
        .page_members(&committed, None, PageSize::new(10).unwrap())
        .await
        .unwrap();
    let page_paths = page
        .items
        .iter()
        .map(ResultMember::path)
        .collect::<Vec<_>>();
    assert_eq!(page_paths, manifest_paths);
}

#[tokio::test]
async fn large_member_integrity_validation_keeps_manifest_allocation_member_bounded() {
    let (Some(pool), Some(target)) = (test_pool().await, target_store()) else {
        return;
    };
    let inventory = ResultInventory::new(pool);
    let attempt = unique_id("attempt");
    let member_id = unique_id("member");
    inventory.create(&attempt).await.unwrap();
    let request = CreateObjectRequest::for_attempt_object(&attempt, &member_id, None);
    let mut body = target.open_create_only(&request).await.unwrap();
    let chunk = vec![0x5a; 1024 * 1024];
    for _ in 0..32 {
        body.write_all(&chunk).await.unwrap();
    }
    let written = body.finish_object().await.unwrap();
    let stored = StoredObject {
        identity: written.identity,
        content_sha256: written.content_sha256,
        size: written.bytes_written,
    };
    let member = ResultMember::new(
        member_id,
        "large/member.bin",
        stored.size,
        stored.identity,
        stored.content_sha256,
    )
    .unwrap();
    inventory.add_member(&attempt, &member).await.unwrap();
    let (committed, metrics) = inventory
        .create_manifest_with_metrics(&attempt, &unique_id("manifest"), &target)
        .await
        .unwrap();
    assert_eq!(committed.total_bytes, 32 * 1024 * 1024);
    assert_eq!(metrics.members_streamed, 1);
    assert!(metrics.max_encoded_member_bytes < 1_024);
}

fn member_identity(attempt: &OpaqueId, object: &OpaqueId) -> ObjectIdentity {
    ObjectIdentity {
        target_store_id: "fixture-store".into(),
        bucket: "fixture-bucket".into(),
        key: format!("_remora/a/{}/o/{}", attempt.as_str(), object.as_str()),
        version_id: None,
    }
}

fn placeholder_publication(attempt: &OpaqueId) -> remora_result::PublicationInput {
    let manifest_id = unique_id("manifest");
    remora_result::PublicationInput {
        attempt_id: attempt.clone(),
        manifest_id: manifest_id.clone(),
        schema_version: MANIFEST_SCHEMA_VERSION,
        manifest: remora_target::StoredObject {
            identity: ObjectIdentity {
                target_store_id: "fixture-store".into(),
                bucket: "fixture-bucket".into(),
                key: remora_target::attempt_manifest_key(attempt, &manifest_id),
                version_id: None,
            },
            content_sha256: ContentSha256::parse(
                "sha256:0000000000000000000000000000000000000000000000000000000000000000",
            )
            .unwrap(),
            size: 0,
        },
        member_count: 1,
        total_bytes: 0,
    }
}