test(fastrun): add more tests for various persistent cache

GitOrigin-RevId: 45c45c8f19

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"

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