Infrastructure for memory usage tests

tests/conftest.py

@@ -0,0 +1,9 @@
+import pytest
+from tests.helpers import MemoryThreshold
+
+
+@pytest.fixture
+def memory_threshold(expected_usage):
+    """Fixture to check peak memory usage against an expected threshold."""
+    with MemoryThreshold(expected_usage) as tracker:
+        yield tracker

tests/helpers.py

@@ -0,0 +1,32 @@
+import tracemalloc
+
+class MemoryThreshold:
+    """
+    Context manager to check peak memory usage against an expected threshold.
+
+    example usage:
+    with MemoryThreshold(expected_usage):
+        # code that should not exceed expected
+    """
+
+    def __init__(self, expected_usage):
+        """
+        Parameters
+        ----------
+        expected_usage : int
+            Expected peak memory usage in bytes
+        """
+        self.expected_usage = expected_usage
+
+    def __enter__(self):
+        tracemalloc.start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        _, peak = tracemalloc.get_traced_memory()
+        tracemalloc.stop()
+
+        assert peak <= self.expected_usage, (
+            "Peak memory usage exceeded expected usage: "
+            f"{peak / 1024:.2f} KB > {self.expected_usage / 1024:.2f} KB"
+        )

tests/outlier_detection/test_median.py

@@ -10,6 +10,7 @@
     _OnDiskMedian,
     nanmedian3D,
 )
+from tests.helpers import MemoryThreshold
 
 
 def test_disk_appendable_array(tmp_path):
@@ -194,3 +195,40 @@ def test_nanmedian3D():
 
     assert med.dtype == np.float32
     assert np.allclose(med, np.nanmedian(cube, axis=0), equal_nan=True)
+
+
+@pytest.mark.parametrize("in_memory", [True, False])
+def test_memory_computer(in_memory, tmp_path):
+    """
+    Analytically calculate how much memory the median computation
+    is supposed to take, then ensure that the implementation
+    stays near that.
+
+    in_memory=True case allocates the following memory:
+    - one cube size
+    - median array == one frame size
+
+    in_memory=False case allocates the following memory:
+    - one buffer size, which by default is the frame size
+    - median array == one frame size
+
+    add a half-frame-size buffer to the expected memory usage in both cases
+    """
+    shp = (20, 500, 500)
+    cube_size = np.dtype("float32").itemsize * shp[0] * shp[1] * shp[2] #bytes
+    frame_size = cube_size / shp[0]
+
+    # calculate expected memory usage
+    if in_memory:
+        expected_mem = cube_size + frame_size*1.5
+    else:
+        expected_mem = frame_size * 2.5
+
+    # compute the median while tracking memory usage
+    with MemoryThreshold(expected_mem):
+        computer = MedianComputer(shp, in_memory=in_memory, tempdir=tmp_path)
+        for i in range(shp[0]):
+            frame = np.full(shp[1:], i, dtype=np.float32)
+            computer.append(frame, i)
+            del frame
+        computer.evaluate()