genegraph-storage 0.69.2

//! #3: zzarr conformance — Python-zarr golden fixtures, Rust round-trips,
//! leading-axis append, and zzarr CSR persistence.
//!
//! Fixtures are generated by `tests/fixtures/zarr/generate_fixtures.py`
//! (Python `zarr` >= 3.0, bytes + zstd codecs, default chunk key encoding).
//! Run with: `cargo test --release --lib -- zzarr`.

use std::path::{Path, PathBuf};

use crate::zzarr::csr::{CsrMeta, read_csr, write_csr};

use super::tmp_dir;

fn fixture(name: &str) -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("tests/fixtures/zarr")
        .join(name)
}

#[tokio::test]
async fn python_f32_1d_fixture_decodes() {
    let arr = crate::zzarr::open(&fixture("f32_1d.zarr")).unwrap();
    assert_eq!(arr.shape(), vec![10]);
    let got = arr.read_all::<f32>().unwrap();
    let want: Vec<f32> = (0..10).map(|i| i as f32 * 0.5).collect();
    assert_eq!(got, want, "Python-written f32 zstd array must decode");
}

#[tokio::test]
async fn python_i64_1d_fixture_decodes() {
    let arr = crate::zzarr::open(&fixture("i64_1d.zarr")).unwrap();
    assert_eq!(arr.shape(), vec![10]);
    let got = arr.read_all::<i64>().unwrap();
    assert_eq!(got, (0..10).collect::<Vec<i64>>());
}

#[tokio::test]
async fn python_f64_2d_multichunk_fixture_decodes_row_major() {
    let arr = crate::zzarr::open(&fixture("f64_2d.zarr")).unwrap();
    assert_eq!(arr.shape(), vec![4, 4]);
    let got = arr.read_all::<f64>().unwrap();
    let want: Vec<f64> = (0..16).map(|i| i as f64 / 10.0).collect();
    assert_eq!(got, want, "row-major layout must be preserved");
}

#[tokio::test]
async fn python_uncompressed_fixture_decodes() {
    let arr = crate::zzarr::open(&fixture("f32_1d_uncompressed.zarr")).unwrap();
    let got = arr.read_all::<f32>().unwrap();
    assert_eq!(got, vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0]);
}

#[tokio::test]
async fn python_f32_1d_subset_read_matches_full_read() {
    let arr = crate::zzarr::open(&fixture("f32_1d.zarr")).unwrap();
    let subset = arr
        .read_subset::<f32>(std::slice::from_ref(&(2..7)))
        .unwrap();
    assert_eq!(subset, vec![1.0f32, 1.5, 2.0, 2.5, 3.0]);
}

#[tokio::test]
async fn roundtrip_f32_zstd_single_chunk() {
    let dir = tmp_dir("zzarr_roundtrip_f32").await.join("a.zarr");
    let values: Vec<f32> = (0..7).map(|i| i as f32 * 1.5).collect();
    crate::zzarr::write_array(&dir, &[7], &[7], &values, true).unwrap();
    let arr = crate::zzarr::open(&dir).unwrap();
    assert_eq!(arr.read_all::<f32>().unwrap(), values);
}

#[tokio::test]
async fn roundtrip_i64_multichunk() {
    let dir = tmp_dir("zzarr_roundtrip_i64").await.join("a.zarr");
    let values: Vec<i64> = (0..100).collect();
    crate::zzarr::write_array(&dir, &[100], &[16], &values, true).unwrap();
    let arr = crate::zzarr::open(&dir).unwrap();
    assert_eq!(arr.read_all::<i64>().unwrap(), values);
}

#[tokio::test]
async fn append_grows_leading_axis_and_preserves_values() {
    let dir = tmp_dir("zzarr_append").await.join("a.zarr");
    crate::zzarr::write_array(&dir, &[5], &[4], &[1.0f32, 2.0, 3.0, 4.0, 5.0], true).unwrap();
    let mut arr = crate::zzarr::open(&dir).unwrap();
    arr.append(&[6.0f32, 7.0, 8.0]).unwrap();
    assert_eq!(arr.shape(), vec![8]);
    assert_eq!(
        arr.read_all::<f32>().unwrap(),
        vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]
    );
    // Reopen from disk: the resize must be persisted.
    let reopened = crate::zzarr::open(&dir).unwrap();
    assert_eq!(reopened.shape(), vec![8]);
}

#[tokio::test]
async fn csr_roundtrip_preserves_values_and_meta() {
    let dir = tmp_dir("zzarr_csr_roundtrip").await.join("main--matrix");
    let meta = CsrMeta {
        nitems: 3,
        nfeatures: 2,
        nclusters: 1,
        csr_shape: [3, 2],
    };
    let data = vec![0.5f32, 1.5, 2.5];
    let indices = vec![0i64, 1, 1];
    let indptr = vec![0i64, 1, 2, 3];
    write_csr(&dir, &data, &indices, &indptr, &meta).unwrap();

    let got = read_csr(&dir).unwrap();
    assert_eq!(got.data, data);
    assert_eq!(got.indices, indices);
    assert_eq!(got.indptr, indptr);
    assert_eq!(got.meta, meta);
}

#[tokio::test]
async fn csr_layout_matches_python_artifact_contract() {
    // Parity with the Python server's _persist_csr: data.zarr, indices.zarr,
    // indptr.zarr plus meta.json under the slug directory.
    let dir = tmp_dir("zzarr_csr_layout").await.join("main--matrix");
    let meta = CsrMeta {
        nitems: 50,
        nfeatures: 4,
        nclusters: 2,
        csr_shape: [50, 4],
    };
    write_csr(&dir, &[1.0f32], &[0i64], &[0i64, 1], &meta).unwrap();
    for name in ["data.zarr", "indices.zarr", "indptr.zarr"] {
        assert!(
            dir.join(name).join("zarr.json").is_file(),
            "{} missing",
            name
        );
    }
    assert!(dir.join("meta.json").is_file());
}

#[tokio::test]
async fn csr_meta_json_keys_match_python_contract() {
    let dir = tmp_dir("zzarr_csr_meta_keys").await.join("d");
    let meta = CsrMeta {
        nitems: 50,
        nfeatures: 4,
        nclusters: 2,
        csr_shape: [50, 4],
    };
    write_csr(&dir, &[1.0f32], &[0i64], &[0i64, 1], &meta).unwrap();
    let text = std::fs::read_to_string(dir.join("meta.json")).unwrap();
    let value: serde_json::Value = serde_json::from_str(&text).unwrap();
    for key in ["nitems", "nfeatures", "nclusters", "csr_shape"] {
        assert!(value.get(key).is_some(), "meta.json missing key {}", key);
    }
    assert_eq!(value["csr_shape"], serde_json::json!([50, 4]));
}

#[tokio::test]
async fn csr_write_leaves_no_atomic_temp_residue() {
    let dir = tmp_dir("zzarr_csr_no_residue").await.join("d");
    let meta = CsrMeta {
        nitems: 1,
        nfeatures: 1,
        nclusters: 1,
        csr_shape: [1, 1],
    };
    write_csr(&dir, &[1.0f32], &[0i64], &[0i64, 1], &meta).unwrap();
    let leftovers: Vec<_> = std::fs::read_dir(&dir)
        .unwrap()
        .map(|e| e.unwrap().file_name().to_string_lossy().into_owned())
        .filter(|n| n.contains(".tmp"))
        .collect();
    assert!(leftovers.is_empty(), "temp residue: {:?}", leftovers);
}

#[tokio::test]
async fn open_directory_without_zarr_json_is_unsupported_format() {
    let dir = tmp_dir("zzarr_open_empty").await.join("not_a_zarr");
    std::fs::create_dir_all(&dir).unwrap();
    let err = crate::zzarr::open(&dir).unwrap_err();
    assert!(matches!(err, crate::StorageError::UnsupportedFormat(_)));
}

#[tokio::test]
async fn write_array_over_existing_zarr_is_rejected() {
    let dir = tmp_dir("zzarr_write_exists").await.join("a.zarr");
    crate::zzarr::write_array(&dir, &[3], &[3], &[1.0f32, 2.0, 3.0], true).unwrap();
    let err = crate::zzarr::write_array(&dir, &[3], &[3], &[1.0f32, 2.0, 3.0], true).unwrap_err();
    assert!(matches!(err, crate::StorageError::InvalidState(_)));
}

/// Writes the arrays the Python verifier reads back
/// (`tests/fixtures/zarr/verify_rust_output.py`). Run explicitly:
/// `cargo test --release --lib -- --ignored zzarr_conformance_output`
#[tokio::test]
#[ignore = "writes conformance output for the Python verifier"]
async fn zzarr_conformance_output_for_python_verifier() {
    let root = fixture("out_rust");
    let _ = std::fs::remove_dir_all(&root);
    let f32_values: Vec<f32> = (0..10).map(|i| i as f32 * 0.5).collect();
    crate::zzarr::write_array(&root.join("f32_1d.zarr"), &[10], &[4], &f32_values, true).unwrap();
    crate::zzarr::write_array(
        &root.join("i64_1d.zarr"),
        &[10],
        &[10],
        &(0..10).collect::<Vec<i64>>(),
        true,
    )
    .unwrap();
    write_csr(
        &root.join("main--matrix"),
        &[0.5f32, 1.5, 2.5],
        &[0i64, 1, 1],
        &[0i64, 1, 2, 3],
        &CsrMeta {
            nitems: 3,
            nfeatures: 2,
            nclusters: 1,
            csr_shape: [3, 2],
        },
    )
    .unwrap();
}