在单元测试期间全局禁用joblib.memory缓存

Question

我使用joblib.Memory模块来缓存几个模块中的一些函数。缓存在模块和类中分别初始化。

模块1:

memory = Memory(location='/cache/')
@memory.cache
def heavy_function(...)
.....

模块2:

memory = Memory(location='/cache/')
@memory.cache
def heavy_function2(...)
.....

模块3:

import Module1
import Module2
def heavy_function3(...)
Module1.heavy_function1(...)
Module1.heavy_function1(...)
.....

现在我有了一个单元测试脚本，我想在单元测试期间全局禁用缓存的使用，以确保所有内容都得到正确计算。这是否可以在不通过手动禁用每个模块的情况下实现Module1.memory.cachedir=None还是不删除cachedir？

我目前的解决方案只是手动修补每个内存调用

单元测试1:

from joblib import Memory
import Module1
Module1.memory = Memory(location=None)
...
unittest.run()

单元测试3:

from joblib import Memory
import Module1 # need to import module 1 just to disable its memory
import Module2 # need to import module 2 just to disable its memory
import Modul3
Module1.memory = Memory(location=None)
Module2.memory = Memory(location=None)
...
unittest.run()

我创建的模块越多，就越需要手动修补内存。我认为可能有更好的解决方案。我在下面提出了一个解决方案。

Answer 1

一个解决方案是在运行测试时设置标志或环境变量。然后在初始化内存之前检查这些标志：

模块1

import os
memflag = os.environ.get('UNITTESTING', False)
memory = Memory(location= None if memflag else '/cache/')
@memory.cache
def heavy_function(...)
.....

单元测试1

os.environ["UNITTESTING"] = '1'
import Module1
.....
unittest.run()
del os.environ["UNITTESTING"]

Answer 2

您可以在导入时对其进行修补，如下所示：

with patch("joblib.Memory") as mock_memory:
mock_memory.return_value.cache = lambda x: x
# then proceed with your tests
def test_whatever():
assert True

mock_memory.return_value.cache = lambda x: x设置joblib。Memory.cache来返回输入(它不会修改它装饰的函数(。

Answer 3

要全局禁用joblib.memory缓存，我会考虑调用register_store_backend，用不起任何作用的DummyStoreBackend覆盖默认的"本地"FileSystemStoreBackend。

如下所示，注意DummyStoreBackend是从Memory的单元测试中复制的，我还没有测试它是否能像预期的一样工作

from joblib.memory import register_store_backend
from joblib._store_backends import StoreBackendBase
class DummyStoreBackend(StoreBackendBase):
"""A dummy store backend that does nothing."""
def _open_item(self, *args, **kwargs):
"""Open an item on store."""
"Does nothing"
def _item_exists(self, location):
"""Check if an item location exists."""
"Does nothing"
def _move_item(self, src, dst):
"""Move an item from src to dst in store."""
"Does nothing"
def create_location(self, location):
"""Create location on store."""
"Does nothing"
def exists(self, obj):
"""Check if an object exists in the store"""
return False
def clear_location(self, obj):
"""Clear object on store"""
"Does nothing"
def get_items(self):
"""Returns the whole list of items available in cache."""
return []
def configure(self, location, *args, **kwargs):
"""Configure the store"""
"Does nothing"
register_store_backend("local", DummyStoreBackend)