Add a new --fail-on-increase option

As requested from some users, add a new --fail-on-increase option that
makes the test run fail if the memory usage increases from previous runs
in tests marked with "limit_memory".

Signed-off-by: Pablo Galindo <[email protected]>

Author: pablogsal

README.md

@@ -97,6 +97,7 @@ MEMORY PROBLEMS demo/test_ok.py::test_memory_exceed
 - `--stacks=STACKS` - Show the N stack entries when showing tracebacks of memory allocations
 - `--native` - Show native frames when showing tracebacks of memory allocations (will be slower)
 - `--trace-python-allocators` - Record allocations made by the Pymalloc allocator (will be slower)
+- `--fail-on-increase` - Fail if the memory usage increases from previous runs in tests marked with "limit_memory"
 
 ## Configuration - INI
 
@@ -105,7 +106,8 @@ MEMORY PROBLEMS demo/test_ok.py::test_memory_exceed
 - `hide_memray_summary(bool)` - hide the memray summary at the end of the execution
 - `stacks(int)` - Show the N stack entries when showing tracebacks of memory allocations
 - `native(bool)`- Show native frames when showing tracebacks of memory allocations (will be slower)
-- `trace_python_allocators` - Record allocations made by the Pymalloc allocator (will be slower)
+- `trace_python_allocators(bool)` - Record allocations made by the Pymalloc allocator (will be slower)
+- `fail-on-increase(bool)` - Fail if the memory usage increases from previous runs in tests marked with "limit_memory"
 
 ## License
 

docs/configuration.rst

@@ -29,6 +29,9 @@ The complete list of command line options is:
 
   ``--trace-python-allocators``
     Record allocations made by the Pymalloc allocator (will be slower)
+  
+  ``--fail-on-increase``
+    Fail if the memory usage increases from previous runs in tests marked with "limit_memory"
 
 .. tab:: Config file options
 
@@ -49,3 +52,6 @@ The complete list of command line options is:
 
   ``trace_python_allocators(bool)``
     Record allocations made by the Pymalloc allocator (will be slower)
+
+  ``--fail-on-increase(bool)``
+    Fail if the memory usage increases from previous runs in tests marked with "limit_memory"

docs/news/91.feature.rst

@@ -0,0 +1 @@
+Add a new --fail-on-increase option that makes the test run fail if the memory usage increases from previous runs in tests marked with "limit_memory"

src/pytest_memray/marks.py

@@ -127,6 +127,28 @@ def long_repr(self) -> str:
         )
 
 
+@dataclass
+class _MoreMemoryInfo:
+    previous_memory: float
+    new_memory: float
+
+    @property
+    def section(self) -> PytestSection:
+        """Return a tuple in the format expected by section reporters."""
+        return (
+            "memray-max-memory",
+            "Test uses more memory than previous run",
+        )
+
+    @property
+    def long_repr(self) -> str:
+        """Generate a longrepr user-facing error message."""
+        return (
+            f"Test previously used {sizeof_fmt(self.previous_memory)} "
+            f"but now uses {sizeof_fmt(self.new_memory)}"
+        )
+
+
 def _generate_section_text(
     allocations: list[AllocationRecord], native_stacks: bool, num_stacks: int
 ) -> str:
@@ -162,15 +184,30 @@ def _passes_filter(
 
 
 def limit_memory(
-    limit: str, *, _result_file: Path, _config: Config
-) -> _MemoryInfo | None:
+    limit: str,
+    *,
+    _result_file: Path,
+    _config: Config,
+    _test_id: str,
+) -> _MemoryInfo | _MoreMemoryInfo | None:
     """Limit memory used by the test."""
     reader = FileReader(_result_file)
     allocations: list[AllocationRecord] = list(
         reader.get_high_watermark_allocation_records(merge_threads=True)
     )
     max_memory = parse_memory_string(limit)
     total_allocated_memory = sum(record.size for record in allocations)
+
+    if _config.cache is not None:
+        cache = _config.cache.get(f"memray/{_test_id}", {})
+        previous = cache.get("total_allocated_memory", float("inf"))
+        fail_on_increase = cast(bool, value_or_ini(_config, "fail_on_increase"))
+        if fail_on_increase and total_allocated_memory > previous:
+            return _MoreMemoryInfo(previous, total_allocated_memory)
+
+        cache["total_allocated_memory"] = total_allocated_memory
+        _config.cache.set(f"memray/{_test_id}", cache)
+
     if total_allocated_memory < max_memory:
         return None
     num_stacks: int = cast(int, value_or_ini(_config, "stacks"))
@@ -190,6 +227,7 @@ def limit_leaks(
     filter_fn: Optional[LeaksFilterFunction] = None,
     _result_file: Path,
     _config: Config,
+    _test_id: str,
 ) -> _LeakedInfo | None:
     reader = FileReader(_result_file)
     allocations: list[AllocationRecord] = list(

src/pytest_memray/plugin.py

@@ -52,6 +52,7 @@ def __call__(
         *args: Any,
         _result_file: Path,
         _config: Config,
+        _test_id: str,
         **kwargs: Any,
     ) -> SectionMetadata | None:
         ...
@@ -237,6 +238,7 @@ def pytest_runtest_makereport(
                 **marker.kwargs,
                 _result_file=result.result_file,
                 _config=self.config,
+                _test_id=item.nodeid,
             )
             if res:
                 report.outcome = "failed"
@@ -382,6 +384,12 @@ def pytest_addoption(parser: Parser) -> None:
         default=False,
         help="Record allocations made by the Pymalloc allocator (will be slower)",
     )
+    group.addoption(
+        "--fail-on-increase",
+        action="store_true",
+        default=False,
+        help="Fail if the memory usage increases from previous runs in tests marked with limit_memory",
+    )
 
     parser.addini("memray", "Activate pytest.ini setting", type="bool")
     parser.addini(
@@ -405,6 +413,11 @@ def pytest_addoption(parser: Parser) -> None:
         help="Record allocations made by the Pymalloc allocator (will be slower)",
         type="bool",
     )
+    parser.addini(
+        "fail-on-increase",
+        help="Fail if the memory usage increases from previous runs in tests marked with limit_memory",
+        type="bool",
+    )
     help_msg = "Show the N tests that allocate most memory (N=0 for all)"
     parser.addini("most_allocations", help_msg)
 

tests/test_pytest_memray.py

@@ -742,3 +742,68 @@ def test_bar():
 
     output = result.stdout.str()
     assert "Only one Memray marker can be applied to each test" in output
+
+
+def test_fail_on_increase(pytester: Pytester):
+    pytester.makepyfile(
+        """
+        import pytest
+        from memray._test import MemoryAllocator
+        allocator = MemoryAllocator()
+
+        @pytest.mark.limit_memory("100MB")
+        def test_memory_alloc_fails():
+            allocator.valloc(1024)
+            allocator.free()
+        """
+    )
+    result = pytester.runpytest("--memray")
+    assert result.ret == ExitCode.OK
+    pytester.makepyfile(
+        """
+        import pytest
+        from memray._test import MemoryAllocator
+        allocator = MemoryAllocator()
+
+        @pytest.mark.limit_memory("100MB")
+        def test_memory_alloc_fails():
+            allocator.valloc(1024 * 10)
+            allocator.free()
+        """
+    )
+    result = pytester.runpytest("--memray", "--fail-on-increase")
+    assert result.ret == ExitCode.TESTS_FAILED
+    output = result.stdout.str()
+    assert "Test uses more memory than previous run" in output
+    assert "Test previously used 1.0KiB but now uses 10.0KiB" in output
+
+
+def test_fail_on_increase_unset(pytester: Pytester):
+    pytester.makepyfile(
+        """
+        import pytest
+        from memray._test import MemoryAllocator
+        allocator = MemoryAllocator()
+
+        @pytest.mark.limit_memory("100MB")
+        def test_memory_alloc_fails():
+            allocator.valloc(1024)
+            allocator.free()
+        """
+    )
+    result = pytester.runpytest("--memray")
+    assert result.ret == ExitCode.OK
+    pytester.makepyfile(
+        """
+        import pytest
+        from memray._test import MemoryAllocator
+        allocator = MemoryAllocator()
+
+        @pytest.mark.limit_memory("100MB")
+        def test_memory_alloc_fails():
+            allocator.valloc(1024 * 10)
+            allocator.free()
+        """
+    )
+    result = pytester.runpytest("--memray")
+    assert result.ret == ExitCode.OK