Skip to content

Commit

Permalink
Add simple round-trip test for npz
Browse files Browse the repository at this point in the history
  • Loading branch information
jturner314 committed Sep 14, 2024
1 parent b7e27b6 commit 9156a07
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 0 deletions.
2 changes: 2 additions & 0 deletions tests/integration/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ use std::io::{self, Read};
use std::ops::{Deref, DerefMut};

mod examples;
#[cfg(feature = "npz")]
mod npz;
mod primitive;
mod round_trip;

Expand Down
51 changes: 51 additions & 0 deletions tests/integration/npz.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
//! .npz examples.

use ndarray::{array, Array2};
use ndarray_npy::{NpzReader, NpzWriter};
use std::{error::Error, io::Cursor};

#[test]
fn round_trip_npz() -> Result<(), Box<dyn Error>> {
let mut buf = Vec::<u8>::new();

let arr1 = array![[1i32, 3, 0], [4, 7, -1]];
let arr2 = array![[9i32, 6], [-5, 2], [3, -1]];

{
let mut writer = NpzWriter::new(Cursor::new(&mut buf));
writer.add_array("arr1", &arr1)?;
writer.add_array("arr2", &arr2)?;
}

{
let mut reader = NpzReader::new(Cursor::new(&buf))?;
assert!(!reader.is_empty());
assert_eq!(reader.len(), 2);
assert_eq!(
reader.names()?,
vec!["arr1".to_string(), "arr2".to_string()],
);
{
let by_name: Array2<i32> = reader.by_name("arr1")?;
assert_eq!(by_name, arr1);
}
{
let by_name: Array2<i32> = reader.by_name("arr1.npy")?;
assert_eq!(by_name, arr1);
}
{
let by_name: Array2<i32> = reader.by_name("arr2")?;
assert_eq!(by_name, arr2);
}
{
let by_name: Array2<i32> = reader.by_name("arr2.npy")?;
assert_eq!(by_name, arr2);
}
{
let res: Result<Array2<i32>, _> = reader.by_name("arr1.npy.npy");
assert!(res.is_err());
}
}

Ok(())
}

0 comments on commit 9156a07

Please sign in to comment.