Search code examples
pythonunit-testingcachingjoblib

Disable joblib.memory caching globally during unittest


I use the joblib.Memory module to cache some functions within several modules. The cache is initialized within modules and classes separately.

Module1:

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

Module2:

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

Module3:

import Module1
import Module2

def heavy_function3(...)
    Module1.heavy_function1(...)
    Module1.heavy_function1(...)
    .....

Now I have a unit testing script and I want to disable the usage of the cache globally during the unit testing to make sure everything is correctly computed. Is this possible without manually disabling it for each module via Module1.memory.cachedir=None or without deleting the cachedir?

My current solution just patches each memory call manually

unittest1:

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

unittest3:

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

The more modules I create, the more manual patching up of the Memory I need. I thought there might be a better solution. One work-around is proposed by me below.


Solution

  • You can patch it at import time like so:

    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 sets joblib.Memory.cache to return the input, ie it won't modify the function that it decorates.