Disable joblib.memory caching globally during unittest

Question:

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.

Asked By: skjerns

||

Answers:

One work-around is to set a flag or an environment variable when running tests. Then check for these flags before initializing the Memory:

Module1

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

unittest1

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

To disable joblib.memory caching globally I would look at calling register_store_backend to overwrite the default ‘local’ FileSystemStoreBackend with a DummyStoreBackend that does nothing.

Something like the following, noting that DummyStoreBackend is copied from the unit tests for Memory, and I have not yet tested if this works as intended

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)
Answered By: gb96

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.

Answered By: Daniel