update tests

Author: ohmayr

tests/unit/utils/test_cache.py

@@ -33,13 +33,13 @@ def bar(self):
     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_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():
@@ -67,16 +67,16 @@ def with_context(self, collisions, *, skip_fields=False, visited_messages=None):
         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.with_context(collisions={"a"}) == "val-1", "a"
+        assert foo.with_context(collisions={"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.with_context(collisions={"b"}) == "val-2"
         assert foo.call_count == 2
 
     # 3. Context Cleared
     # Everything should be forgotten now.
-    assert cache.generation_cache.get() is None
-    assert foo.with_context({"a"}) == "val-7"
+    assert getattr(cache._thread_local, "cache", None) is None
+    assert foo.with_context(collisions={"a"}) == "val-3"