From 9156a07e5f549725ea510a035d24ce120630d9c0 Mon Sep 17 00:00:00 2001
From: Jim Turner <github@turner.link>
Date: Sat, 14 Sep 2024 00:16:16 -0400
Subject: [PATCH] Add simple round-trip test for npz

---
 tests/integration/main.rs |  2 ++
 tests/integration/npz.rs  | 51 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100644 tests/integration/npz.rs

diff --git a/tests/integration/main.rs b/tests/integration/main.rs
index 25a20f4..a8b9197 100644
--- a/tests/integration/main.rs
+++ b/tests/integration/main.rs
@@ -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;
 
diff --git a/tests/integration/npz.rs b/tests/integration/npz.rs
new file mode 100644
index 0000000..349b30e
--- /dev/null
+++ b/tests/integration/npz.rs
@@ -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(())
+}