add test cases for cache

Author: ohmayr

tests/unit/utils/test_cache.py

@@ -31,3 +31,48 @@ def bar(self):
     assert foo.call_count == 1
     assert foo.bar == 42
     assert foo.call_count == 1
+
+
+def test_generation_cache_context():
+    assert cache.generation_cache.get() is None
+    with cache.generation_cache_context():
+        assert isinstance(cache.generation_cache.get(), dict)
+        cache.generation_cache.get()["foo"] = "bar"
+        assert cache.generation_cache.get()["foo"] == "bar"
+    assert cache.generation_cache.get() is None
+
+
+def test_cached_proto_context():
+    class Foo:
+        def __init__(self):
+            self.call_count = 0
+
+        @cache.cached_proto_context
+        def bar(self, *, collisions, other=None):
+            self.call_count += 1
+            return f"baz-{self.call_count}"
+
+    foo = Foo()
+
+    # Without context, no caching
+    assert foo.bar(collisions={"a", "b"}) == "baz-1"
+    assert foo.bar(collisions={"a", "b"}) == "baz-2"
+
+    # With context, caching works
+    with cache.generation_cache_context():
+        assert foo.bar(collisions={"a", "b"}) == "baz-3"
+        assert foo.bar(collisions={"a", "b"}) == "baz-3"
+        assert foo.call_count == 3
+
+        # Different collisions, different cache entry
+        assert foo.bar(collisions={"c"}) == "baz-4"
+        assert foo.bar(collisions={"c"}) == "baz-4"
+        assert foo.call_count == 4
+
+        # Same collisions again, still cached
+        assert foo.bar(collisions={"a", "b"}) == "baz-3"
+        assert foo.call_count == 4
+
+    # Context cleared
+    assert cache.generation_cache.get() is None
+    assert foo.bar(collisions={"a", "b"}) == "baz-5"