From 75bb518657223acee66f201f27afe58a12473609 Mon Sep 17 00:00:00 2001
From: w-bonelli <wesbonelli@gmail.com>
Date: Sat, 12 Nov 2022 08:41:23 -0500
Subject: [PATCH] test(HeadUFile): expand tests for HeadUFile.get_ts()

---
 autotest/test_binaryfile.py | 32 +++++++++++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/autotest/test_binaryfile.py b/autotest/test_binaryfile.py
index 03886d89b5..3ded591e90 100644
--- a/autotest/test_binaryfile.py
+++ b/autotest/test_binaryfile.py
@@ -117,7 +117,7 @@ def test_plot_binary_head_file(example_data_path):
     plt.close()
 
 
-def test_headu_file(tmpdir, example_data_path):
+def test_headu_file_data(tmpdir, example_data_path):
     fname = str(example_data_path / "unstructured" / "headu.githds")
     headobj = HeadUFile(fname)
     assert isinstance(headobj, HeadUFile)
@@ -125,6 +125,7 @@ def test_headu_file(tmpdir, example_data_path):
 
     # ensure recordarray is has correct data
     ra = headobj.recordarray
+    nnodes = 19479
     assert ra["kstp"].min() == 1
     assert ra["kstp"].max() == 1
     assert ra["kper"].min() == 1
@@ -132,7 +133,7 @@ def test_headu_file(tmpdir, example_data_path):
     assert ra["ncol"].min() == 1
     assert ra["ncol"].max() == 14001
     assert ra["nrow"].min() == 7801
-    assert ra["nrow"].max() == 19479
+    assert ra["nrow"].max() == nnodes
 
     # read the heads for the last time and make sure they are correct
     data = headobj.get_data()
@@ -146,7 +147,32 @@ def test_headu_file(tmpdir, example_data_path):
         t1 = np.array([d.min(), d.max()])
         assert np.allclose(t1, minmaxtrue[i])
 
-    return
+
+@pytest.mark.slow
+def test_headufile_get_ts(example_data_path):
+    heads = HeadUFile(example_data_path / "unstructured" / "headu.githds")
+    nnodes = 19479
+
+    # make sure timeseries can be retrieved for each node
+    for i in range(0, nnodes, 100):
+        heads.get_ts(idx=i)
+    with pytest.raises(IndexError):
+        heads.get_ts(idx=i+100)
+
+    # ...and retrieved in groups
+    for i in range(10):
+        heads.get_ts([i, i + 1, i + 2])
+
+    heads = HeadUFile(example_data_path / "mfusg_test" / "01A_nestedgrid_nognc" / "output" / "flow.hds")
+    nnodes = 121
+    for i in range(nnodes):
+        heads.get_ts(idx=i)
+    with pytest.raises(IndexError):
+        heads.get_ts(idx=i+1)
+
+    # ...and retrieved in groups
+    for i in range(10):
+        heads.get_ts([i, i + 1, i + 2])
 
 
 def test_get_headfile_precision(example_data_path):