From 833529a880187981e4a3f8b548275eda485e4a71 Mon Sep 17 00:00:00 2001
From: Andrew Johnson <drewejohnson@users.noreply.github.com>
Date: Fri, 4 Dec 2020 08:28:52 -0800
Subject: [PATCH] Test ResultsList.get_step_where

---
 tests/unit_tests/test_deplete_resultslist.py | 63 ++++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/tests/unit_tests/test_deplete_resultslist.py b/tests/unit_tests/test_deplete_resultslist.py
index 9b40a1ac91..8bfabd3954 100644
--- a/tests/unit_tests/test_deplete_resultslist.py
+++ b/tests/unit_tests/test_deplete_resultslist.py
@@ -65,3 +65,66 @@ def test_get_eigenvalue(res):
     np.testing.assert_allclose(t, t_ref)
     np.testing.assert_allclose(k[:, 0], k_ref)
     np.testing.assert_allclose(k[:, 1], u_ref)
+
+
+@pytest.mark.parametrize("unit", ("s", "d", "min", "h"))
+def test_get_steps(unit):
+    # Make a ResultsList full of near-empty Result instances
+    # Just fill out a time schedule
+    results = openmc.deplete.ResultsList()
+    # Time in units of unit
+    times = np.linspace(0, 100, num=5)
+    if unit == "d":
+        conversion_to_seconds = 60 * 60 * 24
+    elif unit == "h":
+        conversion_to_seconds = 60 * 60
+    elif unit == "min":
+        conversion_to_seconds = 60
+    else:
+        conversion_to_seconds = 1
+
+    for ix in range(times.size):
+        res = openmc.deplete.Results()
+        res.time = times[ix:ix + 1] * conversion_to_seconds
+        results.append(res)
+
+    for expected, value in enumerate(times):
+        actual = results.get_step_where(
+            value, time_units=unit, atol=0, rtol=0)
+        assert actual == expected, (value, results[actual].time[0])
+
+    with pytest.raises(ValueError):
+        # Emulate a result file with a non-zero initial point in time
+        # as in starting from a restart
+        results.get_step_where(times[0] - 1, time_units=unit, atol=0, rtol=0)
+
+    with pytest.raises(ValueError):
+        results.get_step_where(times[-1] + 1, time_units=unit, atol=0, rtol=0)
+
+    # Grab intermediate points with a small offset
+    delta = (times[1] - times[0])
+    offset = delta * 0.1
+    for expected, value in enumerate(times[1:-1], start=1):
+        # Shoot a little low and a little high
+        for mult in {1, -1}:
+            target = value + mult * offset
+            # Compare using absolute and relative tolerances
+            actual = results.get_step_where(
+                target, time_units=unit, atol=offset * 2, rtol=None)
+            assert actual == expected, (
+                target, times[actual], times[expected], offset)
+
+            actual = results.get_step_where(
+                target, time_units=unit, atol=None, rtol=offset / value)
+            assert actual == expected, (
+                target, times[actual], times[expected], offset)
+        # Check that the lower index is returned for the exact mid-point
+        target = value + delta * 0.5
+        actual = results.get_step_where(
+            target, time_units=unit, atol=delta, rtol=delta / value)
+        assert actual == expected
+
+    # Shoot way over with no tolerance -> just give closest value
+    actual = results.get_step_where(
+        times[-1] * 100, time_units=unit, atol=None, rtol=None)
+    assert actual == times.size - 1