From fd24f4875323bfbbd442ba3c35ebde2eb9c11935 Mon Sep 17 00:00:00 2001
From: Matt Wozniski <mwozniski@bloomberg.net>
Date: Fri, 28 Jun 2024 19:21:57 -0400
Subject: [PATCH] Cooperate with anyio

`anyio` allows running `async def` test functions, but the wrapper
installed by Memray to add tracking around the test function breaks
`anyio`'s detection.

Work around this by using an `async def` wrapper when the function being
wrapped is a coroutine function.

Signed-off-by: Matt Wozniski <mwozniski@bloomberg.net>
---
 pyproject.toml              |  1 +
 src/pytest_memray/plugin.py | 79 ++++++++++++++++++++++---------------
 tests/test_pytest_memray.py | 34 ++++++++++++++++
 3 files changed, 82 insertions(+), 32 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index a92ba22..ecee2b5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,6 +37,7 @@ optional-dependencies.lint = [
   "mypy==0.991",
 ]
 optional-dependencies.test = [
+  "anyio>=4.4.0",
   "covdefaults>=2.2.2",
   "pytest>=7.2",
   "coverage>=7.0.5",
diff --git a/src/pytest_memray/plugin.py b/src/pytest_memray/plugin.py
index aa81e65..c72b184 100644
--- a/src/pytest_memray/plugin.py
+++ b/src/pytest_memray/plugin.py
@@ -7,6 +7,7 @@
 import os
 import pickle
 import uuid
+from contextlib import contextmanager
 from dataclasses import dataclass
 from itertools import islice
 from pathlib import Path
@@ -178,39 +179,53 @@ def _build_bin_path() -> Path:
         if markers and "limit_leaks" in markers:
             native = trace_python_allocators = True
 
-        @functools.wraps(func)
-        def wrapper(*args: Any, **kwargs: Any) -> object | None:
-            test_result: object | Any = None
+        @contextmanager
+        def memory_reporting() -> Generator[None, None, None]:
+            # Restore the original function. This is needed because some
+            # pytest plugins (e.g. flaky) will call our pytest_pyfunc_call
+            # hook again with whatever is here, which will cause the wrapper
+            # to be wrapped again.
+            pyfuncitem.obj = func
+
+            result_file = _build_bin_path()
+            with Tracker(
+                result_file,
+                native_traces=native,
+                trace_python_allocators=trace_python_allocators,
+                file_format=FileFormat.AGGREGATED_ALLOCATIONS,
+            ):
+                yield
+
             try:
-                result_file = _build_bin_path()
-                with Tracker(
-                    result_file,
-                    native_traces=native,
-                    trace_python_allocators=trace_python_allocators,
-                    file_format=FileFormat.AGGREGATED_ALLOCATIONS,
-                ):
-                    test_result = func(*args, **kwargs)
-                try:
-                    metadata = FileReader(result_file).metadata
-                except OSError:
-                    return None
-                result = Result(pyfuncitem.nodeid, metadata, result_file)
-                metadata_path = (
-                    self.result_metadata_path
-                    / result_file.with_suffix(".metadata").name
-                )
-                with open(metadata_path, "wb") as file_handler:
-                    pickle.dump(result, file_handler)
-                self.results[pyfuncitem.nodeid] = result
-            finally:
-                # Restore the original function. This is needed because some
-                # pytest plugins (e.g. flaky) will call our pytest_pyfunc_call
-                # hook again with whatever is here, which will cause the wrapper
-                # to be wrapped again.
-                pyfuncitem.obj = func
-            return test_result
-
-        pyfuncitem.obj = wrapper
+                metadata = FileReader(result_file).metadata
+            except OSError:
+                return
+            result = Result(pyfuncitem.nodeid, metadata, result_file)
+            metadata_path = (
+                self.result_metadata_path
+                / result_file.with_suffix(".metadata").name
+            )
+            with open(metadata_path, "wb") as file_handler:
+                pickle.dump(result, file_handler)
+            self.results[pyfuncitem.nodeid] = result
+
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            with memory_reporting():
+                return func(*args, **kwargs)
+
+
+        @functools.wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with memory_reporting():
+                return await func(*args, **kwargs)
+
+        if inspect.iscoroutinefunction(func):
+            pyfuncitem.obj = async_wrapper
+        else:
+            pyfuncitem.obj = wrapper
+
         yield
 
     @hookimpl(hookwrapper=True)
diff --git a/tests/test_pytest_memray.py b/tests/test_pytest_memray.py
index 56d822f..f677e85 100644
--- a/tests/test_pytest_memray.py
+++ b/tests/test_pytest_memray.py
@@ -918,3 +918,37 @@ def test_memory_alloc_fails():
     result = pytester.runpytest("--memray")
 
     assert result.ret == ExitCode.OK
+
+
+def test_running_async_tests_with_anyio(pytester: Pytester) -> None:
+    xml_output_file = pytester.makefile(".xml", "")
+    pytester.makepyfile(
+        """
+        import pytest
+        from memray._test import MemoryAllocator
+        allocator = MemoryAllocator()
+
+        @pytest.fixture
+        def anyio_backend():
+            return 'asyncio'
+
+        @pytest.mark.limit_leaks("5KB")
+        @pytest.mark.anyio
+        async def test_memory_alloc_fails():
+            for _ in range(10):
+                allocator.valloc(1024*10)
+                # No free call here
+        """
+    )
+
+    result = pytester.runpytest("--junit-xml", xml_output_file)
+
+    assert result.ret != ExitCode.OK
+
+    root = ET.parse(str(xml_output_file)).getroot()
+    for testcase in root.iter("testcase"):
+        failure = testcase.find("failure")
+        assert failure.text == (
+            "Test was allowed to leak 5.0KiB per location"
+            " but at least one location leaked more"
+        )