test(HeadUFile): expand tests for HeadUFile.get_ts()

autotest/test_binaryfile.py

@@ -117,22 +117,23 @@ 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)
     assert headobj.nlay == 3
 
     # 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
     assert ra["kper"].max() == 5
     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):