How to have pytest place memory limits on tests?

问题

I am using pytest, but I would like to have a decorator that could set a maximum memory usage per test. Similar to this question which was answered with,

@pytest.mark.timeout(300)
def test_foo():
   pass

I want,

@pytest.mark.maxmem(300)
def test_foo():
   pass

EDIT:

I tried,

>>> import os, psutil
>>> import numpy as np
>>> process = psutil.Process(os.getpid())
>>> process.memory_info().rss/1e9
0.01978368
>>> def f():
...     x = np.arange(int(1e9))
... 
>>> process.memory_info().rss/1e9
0.01982464
>>> f()
>>> process.memory_info().rss/1e9
0.019832832

Which doesn't catch the memory allocation in the function.

回答1:

After learning how to limit the memory used and seeing how much memory is currently used, I wrote a decorator that errors out if the memory increment is too high. It's a bit buggy with setting the limits, but it works well enough for me.

import resource, os, psutil
import numpy

def memory_limit(max_mem):
    def decorator(f):
        def wrapper(*args, **kwargs):
            process = psutil.Process(os.getpid())
            prev_limits = resource.getrlimit(resource.RLIMIT_AS)
            resource.setrlimit(resource.RLIMIT_AS, (process.memory_info().rss + max_mem, -1))
            result = f(*args, **kwargs)
            resource.setrlimit(resource.RLIMIT_AS, prev_limits)
            return result
        return wrapper
    return decorator


@memory_limit(int(16e8))
def allocate(N):
    return numpy.arange(N, dtype='u8')

a = [allocate(int(1e8)) for i in range(10)]

try:
    allocate(int(3e8))
except:
    exit(0)
raise Exception("Should have failed")

At least on my machine, code runs and exits without an error.

来源：https://stackoverflow.com/questions/46327566/how-to-have-pytest-place-memory-limits-on-tests

标签

memory

pytest