update tests

Author: ohmayr

tests/unit/utils/test_cache.py

@@ -43,36 +43,75 @@ def test_generation_cache_context():
 
 
 def test_cached_proto_context():
+    # A dummy object to act as a visited message
+    class MockMessage:
+        pass
+
     class Foo:
         def __init__(self):
             self.call_count = 0
 
+        # We define a signature that matches the real Proto.with_context
+        # to ensure arguments are propagated correctly.
         @cache.cached_proto_context
-        def bar(self, *, collisions, other=None):
+        def with_context(self, collisions, *, skip_fields=False, visited_messages=None):
             self.call_count += 1
-            return f"baz-{self.call_count}"
+            return f"val-{self.call_count}"
 
     foo = Foo()
+    msg_a = MockMessage()
+    msg_b = MockMessage()
 
-    # Without context, no caching
-    assert foo.bar(collisions={"a", "b"}) == "baz-1"
-    assert foo.bar(collisions={"a", "b"}) == "baz-2"
+    # 1. Test Bypass (No Context)
+    # The cache is not active, so every call increments the counter.
+    assert foo.with_context(collisions={"a"}) == "val-1"
+    assert foo.with_context(collisions={"a"}) == "val-2"
 
-    # With context, caching works
+    # 2. Test Context Activation
     with cache.generation_cache_context():
-        assert foo.bar(collisions={"a", "b"}) == "baz-3"
-        assert foo.bar(collisions={"a", "b"}) == "baz-3"
+        # Reset counter to make tracking easier
+        foo.call_count = 0
+
+        # A. Basic Cache Hit
+        assert foo.with_context({"a"}) == "val-1"
+        assert foo.with_context({"a"}) == "val-1"  # Hit
+        assert foo.call_count == 1
+
+        # B. Collision Difference
+        # Changing collisions creates a new key
+        assert foo.with_context({"b"}) == "val-2"
+        assert foo.call_count == 2
+
+        # C. skip_fields Difference (The Critical Fix)
+        # Verify that skip_fields=True is cached separately from the default (False).
+        # This prevents the "Husk" object from poisoning the "Full" object cache.
+        assert foo.with_context({"a"}, skip_fields=True) == "val-3"
+        assert foo.with_context({"a"}, skip_fields=True) == "val-3"  # Hit
         assert foo.call_count == 3
-
-        # Different collisions, different cache entry
-        assert foo.bar(collisions={"c"}) == "baz-4"
-        assert foo.bar(collisions={"c"}) == "baz-4"
+        
+        # Verify the original (skip_fields=False) is still accessible
+        assert foo.with_context({"a"}) == "val-1"  # Hit (Still 1)
+
+        # D. visited_messages Difference
+        # Verify that sets of objects are correctly distinguished.
+        
+        # Call with msg_a
+        assert foo.with_context({"a"}, visited_messages={msg_a}) == "val-4"
+        assert foo.with_context({"a"}, visited_messages={msg_a}) == "val-4" # Hit
         assert foo.call_count == 4
 
-        # Same collisions again, still cached
-        assert foo.bar(collisions={"a", "b"}) == "baz-3"
-        assert foo.call_count == 4
+        # Call with msg_b (Should be a Miss)
+        assert foo.with_context({"a"}, visited_messages={msg_b}) == "val-5"
+        assert foo.call_count == 5
+
+        # E. Set Stability
+        # Verify that the order of items in the set doesn't matter (sets are unordered).
+        # {a, b} should cache-hit with {b, a}
+        assert foo.with_context({"a"}, visited_messages={msg_a, msg_b}) == "val-6"
+        assert foo.with_context({"a"}, visited_messages={msg_b, msg_a}) == "val-6" # Hit
+        assert foo.call_count == 6
 
-    # Context cleared
+    # 3. Context Cleared
+    # Everything should be forgotten now.
     assert cache.generation_cache.get() is None
-    assert foo.bar(collisions={"a", "b"}) == "baz-5"
+    assert foo.with_context({"a"}) == "val-7"