diff --git a/imperative/python/requires-test.txt b/imperative/python/requires-test.txt index 05643464..7b33b1ce 100644 --- a/imperative/python/requires-test.txt +++ b/imperative/python/requires-test.txt @@ -2,3 +2,4 @@ pytest==5.3.0 pytest-sphinx==0.3.1 tensorboardX==2.4 six==1.16.0 +redislite ; platform_system == "Linux" or platform_system == "Darwin" diff --git a/imperative/python/test/unit/utils/test_utils.py b/imperative/python/test/unit/utils/test_utils.py index 0ca81072..f32a3d93 100644 --- a/imperative/python/test/unit/utils/test_utils.py +++ b/imperative/python/test/unit/utils/test_utils.py @@ -1,15 +1,69 @@ +import os +import platform + import pytest -from megengine.utils.persistent_cache import _manager +from megengine.utils.persistent_cache import PersistentCacheOnServer + + +@pytest.mark.parametrize("with_flag", [True, False]) +@pytest.mark.skipif( + platform.system() not in {"Linux", "Darwin"}, + reason="redislite not implemented in windows", +) +def test_persistent_cache_redis(monkeypatch, with_flag): + import redislite + + server = redislite.Redis() + monkeypatch.delenv("MGE_FASTRUN_CACHE_TYPE", raising=False) + monkeypatch.setenv( + "MGE_FASTRUN_CACHE_URL", "redis+socket://{}".format(server.socket_file) + ) + if with_flag: + server.set("mgb-cache-flag", 1) + pc = PersistentCacheOnServer() + pc.put("test", "hello", "world") + if with_flag: + pc = PersistentCacheOnServer() + assert pc.get("test", "hello") == b"world" + assert pc.config.type == "redis" + else: + assert pc.config.type == "in-file" + + +def test_persistent_cache_file(monkeypatch, tmp_path): + monkeypatch.setenv("MGE_FASTRUN_CACHE_TYPE", "FILE") + monkeypatch.setenv("MGE_FASTRUN_CACHE_DIR", tmp_path) + pc = PersistentCacheOnServer() + pc.put("test", "store", "this") + assert pc.config.type == "in-file" + del pc + pc = PersistentCacheOnServer() + assert pc.get("test", "store") == b"this" + + +def test_persistent_cache_file_clear(monkeypatch, tmp_path): + monkeypatch.setenv("MGE_FASTRUN_CACHE_TYPE", "FILE") + monkeypatch.setenv("MGE_FASTRUN_CACHE_DIR", tmp_path) + pc = PersistentCacheOnServer() + pc_dummy = PersistentCacheOnServer() + pc.put("test", "drop", "this") + assert pc.config.type == "in-file" + del pc + # this dummy instance shouldn't override cache file + del pc_dummy + os.unlink(os.path.join(tmp_path, "cache.bin")) + pc = PersistentCacheOnServer() + assert pc.get("test", "drop") is None -def test_persistent_cache(): - pc = _manager - k0 = b"\x00\x00" - k1 = b"\x00\x01" - cat = "test" - pc.put(cat, k0, k1) - pc.put(cat, k1, k0) - assert k1 == pc.get(cat, k0) - assert k0 == pc.get(cat, k1) - assert pc.get("test1", k0) == None +def test_persistent_cache_memory(monkeypatch): + monkeypatch.setenv("MGE_FASTRUN_CACHE_TYPE", "MEMORY") + pc = PersistentCacheOnServer() + assert pc.config is None + pc.put("test", "drop", "this") + assert pc.config.type == "in-memory" + assert pc.get("test", "drop") == b"this" + del pc + pc = PersistentCacheOnServer() + assert pc.get("test", "drop") is None