|
1 | 1 | import asyncio |
2 | 2 | import logging |
3 | 3 | import os |
| 4 | +import warnings |
4 | 5 | from math import inf |
5 | 6 | from typing import Any, Callable, no_type_check |
6 | 7 | from unittest.mock import MagicMock |
|
22 | 23 | # Windows |
23 | 24 | resource = None # type:ignore |
24 | 25 |
|
| 26 | +try: |
| 27 | + import tracemalloc |
| 28 | +except ModuleNotFoundError: |
| 29 | + tracemalloc = None |
| 30 | + |
25 | 31 |
|
26 | 32 | @pytest.fixture() |
27 | 33 | def anyio_backend(): |
@@ -207,3 +213,38 @@ async def ipkernel(anyio_backend): |
207 | 213 | yield kernel |
208 | 214 | kernel.destroy() |
209 | 215 | ZMQInteractiveShell.clear_instance() |
| 216 | + |
| 217 | + |
| 218 | +@pytest.fixture() |
| 219 | +def tracemalloc_resource_warning(recwarn, N=10): |
| 220 | + """fixture to enable tracemalloc for a single test, and report the |
| 221 | + location of the leaked resource |
| 222 | +
|
| 223 | + We cannot only enable tracemalloc, as otherwise it is stopped just after the |
| 224 | + test, the frame cache is cleared by tracemalloc.stop() and thus the warning |
| 225 | + printing code get None when doing |
| 226 | + `tracemalloc.get_object_traceback(r.source)`. |
| 227 | +
|
| 228 | + So we need to both filter the warnings to enable ResourceWarning, and loop |
| 229 | + through it print the stack before we stop tracemalloc and continue. |
| 230 | +
|
| 231 | + """ |
| 232 | + if tracemalloc is None: |
| 233 | + yield |
| 234 | + return |
| 235 | + |
| 236 | + tracemalloc.start(N) |
| 237 | + with warnings.catch_warnings(): |
| 238 | + warnings.simplefilter("always", category=ResourceWarning) |
| 239 | + yield None |
| 240 | + try: |
| 241 | + for r in recwarn: |
| 242 | + if r.category is ResourceWarning and r.source is not None: |
| 243 | + tb = tracemalloc.get_object_traceback(r.source) |
| 244 | + if tb: |
| 245 | + info = f"Leaking resource:{r}\n |" + "\n |".join(tb.format()) |
| 246 | + # technically an Error and not a failure as we fail in the fixture |
| 247 | + # and not the test |
| 248 | + pytest.fail(info) |
| 249 | + finally: |
| 250 | + tracemalloc.stop() |
0 commit comments