Question

我正在使用pytest，但我希望有一个可以设置每个测试的最大内存使用量的装饰器。与回答的question类似，

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

我想，

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

编辑：

我试过了，

>>> 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

这不会捕获函数中的内存分配。

Answer 1

在学习了如何limit the memory used并查看了多少memory is currently used之后，我写了一个装饰器，如果内存增量过高则会出错。设置限制有点麻烦，但对我来说效果还不错。

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")

至少在我的机器上，代码运行并退出而没有错误。

如何让pytest在测试中设置内存限制？

1 个答案: