diff --git a/docs/usage.rst b/docs/usage.rst index e27c3e5..dd61ef1 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -40,6 +40,10 @@ validations on tests when this plugin is enabled. ``limit_memory`` ---------------- +.. py:function:: limit_memory(memory_limit: str) + + Fail the execution of the test if the test allocates more memory than allowed + When this marker is applied to a test, it will cause the test to fail if the execution of the test allocates more memory than allowed. It takes a single argument with a string indicating the maximum memory that the test can allocate. @@ -61,3 +65,58 @@ Example of usage: @pytest.mark.limit_memory("24 MB") def test_foobar(): pass # do some stuff that allocates memory + +``check_leaks`` +--------------- + +.. py:function:: check_leaks(max_leaked_memory: str , *, warmups: int=0) + + Fail the execution of the test if the test leaks more memory than allowed. + +.. warning:: + + Running this marker can cause the test to take extra time to run, as it will + activate some additional check in ``memray`` to ensure that all leaks are + detected. Additionally, if the ``warmups`` argument is provided, the test + will be run multiple times before the final execution, which will also + increase the total time it takes to complete the test run. + +When this marker is applied to a test, it will cause the test to fail if the execution +of the test leaks more memory than allowed. It takes a single positional argument with a +string indicating the maximum memory that the test is allowed to leak. + +The format for the string is `` ([KMGTP]B|B)``. The marker will raise +``ValueError`` if the string format cannot be parsed correctly. + +The marker also takes an optional keyword-only argument ``warmups``. This argument +indicates the number of times the test will be executed before the memory leak check +is performed. This is useful to avoid false positives when the test is allocating +memory in caches that will survive the test execution or to account for memory +allocated (and not released) by the interpreter under normal execution. + +.. important:: + + The Python interpreter allocates memory in some internal caches that may not be + released when the test ends. This means that tests using this marker may need to + give some room to account for this or to use the ``warmups`` argument. The amount + of memory used by the interpreter and the location assigned to it can also change + between minor versions of Python. + +Using this marker with "0B" as the maximum allowed leaked memory can be a bit +challenging and will surely require a nonzero value for the ``warmups`` +argument. + +If you need to further investigate the leaks, you can use the ``--memray-bin-path`` +to specify a directory where ``memray`` will store the binary files with the results. +This will allow you to use the ``memray`` CLI to inspect the results using the different +reporters available. Check `the memray docs +`_ for more information. + +Example of usage: + +.. code-block:: python + + @pytest.mark.check_leaks("1 MB", warmups=3) + def test_foobar(): + pass # do some stuff that cannot leak memory + diff --git a/src/pytest_memray/marks.py b/src/pytest_memray/marks.py index 6efbaf0..9c7ab9d 100644 --- a/src/pytest_memray/marks.py +++ b/src/pytest_memray/marks.py @@ -1,10 +1,15 @@ from __future__ import annotations from dataclasses import dataclass +from pathlib import Path +from typing import Any +from typing import Callable from typing import Optional from typing import Tuple from memray import AllocationRecord +from memray import FileReader +from memray import Tracker from .utils import parse_memory_string from .utils import sizeof_fmt @@ -46,15 +51,84 @@ def long_repr(self) -> str: return f"Test was limited to {max_memory_str} but allocated {total_memory_str}" -def limit_memory( - limit: str, *, _allocations: list[AllocationRecord] -) -> Optional[_MemoryInfo]: +@dataclass +class _LeakedInfo: + """Type that holds all memray-related info for a failed test.""" + + max_memory: float + total_allocated_memory: int + allocations: list[AllocationRecord] + + @property + def section(self) -> PytestSection: + """Return a tuple in the format expected by section reporters.""" + total_memory_str = sizeof_fmt(self.total_allocated_memory) + max_memory_str = sizeof_fmt(self.max_memory) + text_lines = [ + f"Test leaked {total_memory_str} out of limit of {max_memory_str}", + "List of leaked allocations: ", + ] + for record in self.allocations: + size = record.size + stack_trace = record.stack_trace() + if not stack_trace: + continue + (function, file, line), *_ = stack_trace + if "pytest_memray" in file: + # Do not report leaks that happen originated from the plugin + continue + text_lines.append(f"\t- {function}:{file}:{line} -> {sizeof_fmt(size)}") + return "memray-leaked-memory", "\n".join(text_lines) + + @property + def long_repr(self) -> str: + """Generate a longrepr user-facing error message.""" + total_memory_str = sizeof_fmt(self.total_allocated_memory) + max_memory_str = sizeof_fmt(self.max_memory) + return ( + f"Test was allowed to leak {max_memory_str} but leaked {total_memory_str}" + ) + + +def limit_memory(limit: str, *, results_file: Path) -> Optional[_MemoryInfo]: """Limit memory used by the test.""" + reader = FileReader(results_file) + allocations = 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) + total_allocated_memory = sum(record.size for record in allocations) if total_allocated_memory < max_memory: return None - return _MemoryInfo(max_memory, total_allocated_memory, _allocations) + return _MemoryInfo(max_memory, total_allocated_memory, allocations) + + +def check_leaks_runner( + func: Callable[[], Any], + results_file: Path, + *marker_args: Any, + warmups: int = 1, + **marker_kwargs: Any, +) -> Any: + result = None + for _ in range(warmups): + with Tracker("/dev/null"): + result = func() + with Tracker(results_file, trace_python_allocators=True): + result = func() + return result + + +def check_leaks( + limit: str, *, results_file: Path, **kwargs: Any +) -> Optional[_LeakedInfo]: + reader = FileReader(results_file) + allocations = list((reader.get_leaked_allocation_records(merge_threads=True))) + max_memory = parse_memory_string(limit) + total_leaked_memory = sum(record.size for record in allocations) + if total_leaked_memory <= max_memory: + return None + return _LeakedInfo(max_memory, total_leaked_memory, allocations) __all__ = [ diff --git a/src/pytest_memray/plugin.py b/src/pytest_memray/plugin.py index d7f2dea..da11147 100644 --- a/src/pytest_memray/plugin.py +++ b/src/pytest_memray/plugin.py @@ -12,9 +12,13 @@ from pathlib import Path from tempfile import TemporaryDirectory from typing import Any +from typing import Callable +from typing import Dict from typing import Generator from typing import Iterable from typing import List +from typing import Mapping +from typing import Protocol from typing import Tuple from typing import cast @@ -33,11 +37,27 @@ from pytest import TestReport from pytest import hookimpl +from .marks import check_leaks +from .marks import check_leaks_runner from .marks import limit_memory from .utils import WriteEnabledDirectoryAction from .utils import sizeof_fmt -MARKERS = {"limit_memory": limit_memory} +TestFunction = Callable[[], Any] + + +class Runner(Protocol): + def __call__( + self, func: TestFunction, results_file: Path, *args: Any, **kwargs: Any + ) -> Any: + ... + + +MARKERS: Dict[str, Callable[..., Any]] = { + "limit_memory": limit_memory, + "check_leaks": check_leaks, +} +MARKERS_RUNNERS: Dict[str, Runner] = {"check_leaks": check_leaks_runner} N_TOP_ALLOCS = 5 N_HISTOGRAM_BINS = 5 @@ -86,7 +106,9 @@ class Manager: def __init__(self, config: Config) -> None: self.results: dict[str, Result] = {} self.config = config - path: Path | None = config.getvalue("memray_bin_path") + the_path = config.getvalue("memray_bin_path") + assert the_path is None or isinstance(the_path, Path) + path: Path | None = the_path self._tmp_dir: None | TemporaryDirectory[str] = None if path is None: # Check the MEMRAY_RESULT_PAtH environment variable. If this @@ -118,10 +140,28 @@ def pytest_unconfigure(self, config: Config) -> Generator[None, None, None]: if os.environ.get("MEMRAY_RESULT_PATH"): del os.environ["MEMRAY_RESULT_PATH"] + def _default_runner( + self, func: Callable[[], Any], results_file: Path, *args: Any, **kwargs: Any + ) -> Any: + result = None + with Tracker(results_file): + result = func() + return result + @hookimpl(hookwrapper=True) # type: ignore[misc] # Untyped decorator - def pytest_pyfunc_call(self, pyfuncitem: Function) -> object | None: + def pytest_pyfunc_call(self, pyfuncitem: Function) -> Generator[None, None, None]: func = pyfuncitem.obj + runner: Runner = self._default_runner + marker_args: Tuple[Any, ...] = tuple() + marker_kwargs: Mapping[str, Any] = dict() + for marker in pyfuncitem.iter_markers(): + if marker.name in MARKERS_RUNNERS: + runner = MARKERS_RUNNERS[marker.name] + marker_args = marker.args + marker_kwargs = marker.kwargs + break + def _build_bin_path() -> Path: if self._tmp_dir is None: of_id = pyfuncitem.nodeid.replace("::", "-") @@ -138,8 +178,10 @@ def _build_bin_path() -> Path: def wrapper(*args: Any, **kwargs: Any) -> object | None: try: result_file = _build_bin_path() - with Tracker(result_file): - result: object | None = func(*args, **kwargs) + runner_function = functools.partial(func, *args, **kwargs) + result: object | None = runner( + runner_function, result_file, *marker_args, **marker_kwargs + ) try: metadata = FileReader(result_file).metadata except OSError: @@ -182,10 +224,9 @@ def pytest_runtest_makereport( result = self.results.get(item.nodeid) if not result: continue - reader = FileReader(result.result_file) - func = reader.get_high_watermark_allocation_records - allocations = list((func(merge_threads=True))) - res = marker_fn(*marker.args, **marker.kwargs, _allocations=allocations) + res = marker_fn( + *marker.args, **marker.kwargs, results_file=result.result_file + ) if res: report.outcome = "failed" report.longrepr = res.long_repr diff --git a/tests/test_pytest_memray.py b/tests/test_pytest_memray.py index 5d2f075..93545fa 100644 --- a/tests/test_pytest_memray.py +++ b/tests/test_pytest_memray.py @@ -469,3 +469,69 @@ def test_memory_alloc_fails_2(): result = pytester.runpytest("--memray", "-n", "2") assert result.ret == outcome + + +@pytest.mark.parametrize( + "size, outcome", + [ + (1024 * 1, ExitCode.OK), + (1024 * 1 + 1, ExitCode.TESTS_FAILED), + (1024 * 2, ExitCode.TESTS_FAILED), + (1024 * 1 - 1, ExitCode.OK), + (1024 * 0.5, ExitCode.OK), + ], +) +def test_leak_marker(pytester: Pytester, size: int, outcome: ExitCode) -> None: + pytester.makepyfile( + f""" + import pytest + from memray._test import MemoryAllocator + allocator = MemoryAllocator() + + @pytest.mark.check_leaks("1KB", warmups=1) + def test_memory_alloc_fails(): + allocator.valloc({size}) + # No free call here + """ + ) + + result = pytester.runpytest("--memray") + + assert result.ret == outcome + + +def test_leak_marker_warmups(pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + + @pytest.mark.check_leaks("1MB", warmups=10) + def test_the_thing(): + pass + """ + ) + + with patch("pytest_memray.marks.Tracker") as mock: + result = pytester.runpytest("--memray") + + assert result.ret == ExitCode.OK + assert mock.call_count == 10 + 1 + + +def test_leak_marker_does_not_work_if_memray_inactive(pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + from memray._test import MemoryAllocator + allocator = MemoryAllocator() + + @pytest.mark.check_leaks("0B", warmups=1) + def test_memory_alloc_fails(): + allocator.valloc(512) + # No free call here + """ + ) + + result = pytester.runpytest("") + + assert result.ret == ExitCode.OK