diff --git a/api/onnx_web/server/model_cache.py b/api/onnx_web/server/model_cache.py index c9f9df4e..b7207c41 100644 --- a/api/onnx_web/server/model_cache.py +++ b/api/onnx_web/server/model_cache.py @@ -53,6 +53,11 @@ class ModelCache: cache.append((tag, key, value)) self.prune() + def clear(self): + global cache + + cache.clear() + def prune(self): global cache diff --git a/api/tests/server/test_model_cache.py b/api/tests/server/test_model_cache.py index 149b22e9..000065d0 100644 --- a/api/tests/server/test_model_cache.py +++ b/api/tests/server/test_model_cache.py @@ -5,18 +5,21 @@ from onnx_web.server.model_cache import ModelCache class TestStringMethods(unittest.TestCase): def test_drop_existing(self): cache = ModelCache(10) + cache.clear() cache.set("foo", ("bar",), {}) self.assertGreater(cache.size, 0) self.assertEqual(cache.drop("foo", ("bar",)), 1) def test_drop_missing(self): cache = ModelCache(10) + cache.clear() cache.set("foo", ("bar",), {}) self.assertGreater(cache.size, 0) self.assertEqual(cache.drop("foo", ("bin",)), 0) def test_get_existing(self): cache = ModelCache(10) + cache.clear() value = {} cache.set("foo", ("bar",), value) self.assertGreater(cache.size, 0) @@ -24,6 +27,7 @@ class TestStringMethods(unittest.TestCase): def test_get_missing(self): cache = ModelCache(10) + cache.clear() value = {} cache.set("foo", ("bar",), value) self.assertGreater(cache.size, 0)