From 07b864f6e2b6ebead7a5440ad95b1ade42e64409 Mon Sep 17 00:00:00 2001 From: Pablo Galindo Date: Wed, 17 Apr 2024 15:50:01 +0100 Subject: [PATCH 1/3] Add a new current_thread_only to all markers Add a new "current_thread_only" keyword to the "limit_memory" and "limit_leaks" markers to ignore all allocations made in threads other than the one running the test. --- docs/news/117.feature.rst | 3 +++ docs/usage.rst | 10 +++++++- pyproject.toml | 2 +- src/pytest_memray/marks.py | 32 +++++++++++++++++++----- tests/test_pytest_memray.py | 50 +++++++++++++++++++++++++++++++++++++ 5 files changed, 89 insertions(+), 8 deletions(-) create mode 100644 docs/news/117.feature.rst diff --git a/docs/news/117.feature.rst b/docs/news/117.feature.rst new file mode 100644 index 0000000..37db920 --- /dev/null +++ b/docs/news/117.feature.rst @@ -0,0 +1,3 @@ +Add a new ``current_thread_only`` keyword argument to the ``limit_memory`` and +``limit_leaks`` markers to ignore all allocations made in threads other than +the one running the test. diff --git a/docs/usage.rst b/docs/usage.rst index 3850d68..e57f03f 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -35,7 +35,7 @@ This plugin provides `markers ([KMGTP]B|B)``. The marker will raise ``ValueError`` if the string format cannot be parsed correctly. + If the optional keyword-only argument ``current_thread_only`` is set to *True*, the + plugin will only track memory allocations made by the current thread and all other + allocations will be ignored. + .. warning:: As the Python interpreter has its own @@ -96,6 +100,10 @@ that can be used to enforce additional checks and validations on tests. ignored, the test will not fail. This can be used to discard any known false positives. + If the optional keyword-only argument ``current_thread_only`` is set to *True*, the + plugin will only track memory allocations made by the current thread and all other + allocations will be ignored. + .. tip:: You can pass the ``--memray-bin-path`` argument to ``pytest`` to specify diff --git a/pyproject.toml b/pyproject.toml index 8708c3a..a92ba22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,7 @@ maintainers = [ requires-python = ">=3.8" dependencies = [ "pytest>=7.2", - "memray>=1.5", + "memray>=1.12", ] optional-dependencies.docs = [ "furo>=2022.12.7", diff --git a/src/pytest_memray/marks.py b/src/pytest_memray/marks.py index 485cf05..eac7e04 100644 --- a/src/pytest_memray/marks.py +++ b/src/pytest_memray/marks.py @@ -186,15 +186,27 @@ def _passes_filter( def limit_memory( limit: str, *, + current_thread_only: bool = False, _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) - ) + allocations: list[AllocationRecord] = [] + if current_thread_only: + main_thread = reader.metadata.main_thread_id + allocations.extend( + record + for record in reader.get_high_watermark_allocation_records( + merge_threads=False + ) + if record.tid == main_thread + ) + else: + allocations.extend( + 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) @@ -225,14 +237,22 @@ def limit_leaks( location_limit: str, *, filter_fn: Optional[LeaksFilterFunction] = None, + current_thread_only: bool = False, _result_file: Path, _config: Config, _test_id: str, ) -> _LeakedInfo | None: reader = FileReader(_result_file) - allocations: list[AllocationRecord] = list( - reader.get_leaked_allocation_records(merge_threads=True) - ) + allocations: list[AllocationRecord] = [] + if current_thread_only: + main_thread_id = reader.metadata.main_thread_id + allocations.extend( + record + for record in reader.get_leaked_allocation_records(merge_threads=False) + if record.tid == main_thread_id + ) + else: + allocations.extend(reader.get_leaked_allocation_records(merge_threads=True)) memory_limit = parse_memory_string(location_limit) diff --git a/tests/test_pytest_memray.py b/tests/test_pytest_memray.py index 7eb4a0b..56d822f 100644 --- a/tests/test_pytest_memray.py +++ b/tests/test_pytest_memray.py @@ -868,3 +868,53 @@ def test_memory_alloc_fails(): ) result = pytester.runpytest("--memray") assert result.ret == ExitCode.OK + + +def test_limit_memory_in_current_thread(pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + from memray._test import MemoryAllocator + allocator = MemoryAllocator() + import threading + def allocating_func(): + for _ in range(10): + allocator.valloc(1024*5) + # No free call here + + @pytest.mark.limit_memory("5KB", current_thread_only=True) + def test_memory_alloc_fails(): + t = threading.Thread(target=allocating_func) + t.start() + t.join() + """ + ) + + result = pytester.runpytest("--memray") + + assert result.ret == ExitCode.OK + + +def test_leaks_in_current_thread(pytester: Pytester) -> None: + pytester.makepyfile( + """ + import pytest + from memray._test import MemoryAllocator + allocator = MemoryAllocator() + import threading + def allocating_func(): + for _ in range(10): + allocator.valloc(1024*5) + # No free call here + + @pytest.mark.limit_leaks("5KB", current_thread_only=True) + def test_memory_alloc_fails(): + t = threading.Thread(target=allocating_func) + t.start() + t.join() + """ + ) + + result = pytester.runpytest("--memray") + + assert result.ret == ExitCode.OK From 3631eed9ee66a4cc1058c244cac2b8bab99d30d6 Mon Sep 17 00:00:00 2001 From: Matt Wozniski Date: Wed, 17 Apr 2024 15:15:13 -0400 Subject: [PATCH 2/3] fixup! Add a new current_thread_only to all markers --- docs/usage.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/usage.rst b/docs/usage.rst index e57f03f..a937c2c 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -67,7 +67,7 @@ that can be used to enforce additional checks and validations on tests. pass # do some stuff that allocates memory -.. py:function:: pytest.mark.limit_leaks(location_limit: str, filter_fn: LeaksFilterFunction | None = None) +.. py:function:: pytest.mark.limit_leaks(location_limit: str, filter_fn: LeaksFilterFunction | None = None, current_thread_only: bool = False) Fail the execution of the test if any call stack in the test leaks more memory than allowed. From f06157b9027cc2fe9809775bff4663d91cd54848 Mon Sep 17 00:00:00 2001 From: Pablo Galindo Salgado Date: Thu, 18 Apr 2024 12:34:07 +0100 Subject: [PATCH 3/3] Apply suggestions from code review Co-authored-by: Matt Wozniski --- src/pytest_memray/marks.py | 34 ++++++++++++---------------------- 1 file changed, 12 insertions(+), 22 deletions(-) diff --git a/src/pytest_memray/marks.py b/src/pytest_memray/marks.py index eac7e04..01fdb49 100644 --- a/src/pytest_memray/marks.py +++ b/src/pytest_memray/marks.py @@ -193,20 +193,13 @@ def limit_memory( ) -> _MemoryInfo | _MoreMemoryInfo | None: """Limit memory used by the test.""" reader = FileReader(_result_file) - allocations: list[AllocationRecord] = [] - if current_thread_only: - main_thread = reader.metadata.main_thread_id - allocations.extend( - record - for record in reader.get_high_watermark_allocation_records( - merge_threads=False - ) - if record.tid == main_thread - ) - else: - allocations.extend( - reader.get_high_watermark_allocation_records(merge_threads=True) + allocations: list[AllocationRecord] = [ + record + for record in reader.get_high_watermark_allocation_records( + merge_threads=not current_thread_only ) + if not current_thread_only or record.tid == reader.metadata.main_thread_id + ] max_memory = parse_memory_string(limit) total_allocated_memory = sum(record.size for record in allocations) @@ -243,16 +236,13 @@ def limit_leaks( _test_id: str, ) -> _LeakedInfo | None: reader = FileReader(_result_file) - allocations: list[AllocationRecord] = [] - if current_thread_only: - main_thread_id = reader.metadata.main_thread_id - allocations.extend( - record - for record in reader.get_leaked_allocation_records(merge_threads=False) - if record.tid == main_thread_id + allocations: list[AllocationRecord] = [ + record + for record in reader.get_leaked_allocation_records( + merge_threads=not current_thread_only ) - else: - allocations.extend(reader.get_leaked_allocation_records(merge_threads=True)) + if not current_thread_only or record.tid == reader.metadata.main_thread_id + ] memory_limit = parse_memory_string(location_limit)