Repository: bloomberg/pytest-memray (language: python)

Issue:
pytest-memray breaks anyio
Hola @pablogsal,

I am facing the same issue: the async tests are skipped if we pass `--memray` argument to pytest.

## Steps to reproduce the issue:
Use the following test file: `test_async.py`
```python
import pytest


@pytest.fixture
def anyio_backend():
    return 'asyncio'

@pytest.mark.anyio
async def test_async():
    assert True
```
Install required dependencies:
```shell
python -m pip install pytest anyio pytest-memray
```

The test runs as expected if `--memray` is not passed:
```shell
 python -m pytest -vv -x test_async.py
```
Output:
```
plugins: memray-1.6.0, anyio-4.0.0
collected 1 item                                                                                                                                                                                          

test_async.py::test_async PASSED
```

However, the test is skipped if we pass `--memray`:
```shell
 python -m pytest --memray -vv -x test_async.py
```
Output:
```
plugins: memray-1.6.0, anyio-4.0.0
collected 1 item                                                                                                                                                                                          

test_async.py::test_async SKIPPED (async def function and no async plugin installed (see warnings))                                                                                                 [100%]

============================================================================================ warnings summary =============================================================================================
test_async.py::test_async
  <MY_PROJECT_PATH>/.venv/lib/python3.9/site-packages/_pytest/python.py:151: PytestUnhandledCoroutineWarning: async def functions are not natively supported and have been skipped.
  You need to install a suitable plugin for your async framework, for example:
    - anyio
    - pytest-asyncio
    - pytest-tornasync
    - pytest-trio
    - pytest-twisted
    warnings.warn(PytestUnhandledCoroutineWarning(msg.format(nodeid)))

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
```

_Originally posted by @albertvillanova in https://github.com/bloomberg/pytest-memray/discussions/101#discussioncomment-9738673_


Discussion / hints:
It looks `anyio` [stores some state on the test function](https://github.com/agronholm/anyio/blob/7b4f94ea2812edc85f388afa271bab4f15b8b608/src/anyio/pytest_plugin.py#L96), and `pytest-memray` replaces the test function with a wrapper: https://github.com/bloomberg/pytest-memray/blob/0654e6b8c57cd98314c27484c280cb54c1b3bb94/src/pytest_memray/plugin.py#L213 Those two things don't play nicely with each other.
Actually, looks like [pytest is storing that state](https://github.com/pytest-dev/pytest/blob/ff7598013581b5820dd83253df7d0cc46f785020/src/_pytest/mark/structures.py#L414) on the test function. `anyio` is just retrieving it.

We probably need to copy the test function's `__dict__` over onto our wrapper.
@godlygeek thanks for the investigation.
My original analysis didn't spot the problem, actually. It's not that anyio isn't finding the state it expect on the wrapper, it's that anyio is completely ignoring the wrapper, because it only wants to process coroutine functions and ignore other stuff, and our wrapper isn't a coroutine function (even though calling it does return a coroutine, anyio is specifically looking only for `async def` functions).

We might be able to get around that by defining an `async def` wrapper if the function we're wrapping is an `async def` coroutine...

Fix patch:
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",


Test patch:
diff --git a/src/pytest_memray/plugin.py b/src/pytest_memray/plugin.py
index aa81e65..6d64428 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,50 @@ 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"
+        )