diff --git a/tests/fixtures/treescope_examples_fixture.py b/tests/fixtures/treescope_examples_fixture.py
index 8000080..3ab249e 100644
--- a/tests/fixtures/treescope_examples_fixture.py
+++ b/tests/fixtures/treescope_examples_fixture.py
@@ -132,6 +132,11 @@ def __call__(self, value: int) -> int:
 some_callable_block = SomethingCallable()
 
 
+@dataclasses.dataclass(frozen=True)
+class FrozenDataclassKey:
+  name: str
+
+
 @jax.tree_util.register_pytree_with_keys_class
 class UnknownPytreeNode:
   """A Pytree node treescope doesn't know."""
@@ -145,7 +150,7 @@ def __repr__(self):
 
   def tree_flatten_with_keys(self):
     return (
-        ((jax.tree_util.GetAttrKey("x"), self.x), ("custom_key", self.y)),
+        ((FrozenDataclassKey("x"), self.x), ("string_key", self.y)),
         "example_pytree_aux_data",
     )
 
diff --git a/tests/renderer_test.py b/tests/renderer_test.py
index fba50d0..7fe7c9f 100644
--- a/tests/renderer_test.py
+++ b/tests/renderer_test.py
@@ -18,6 +18,7 @@
 import collections
 import dataclasses
 import functools
+import re
 import textwrap
 import types
 from typing import Any, Callable
@@ -553,11 +554,16 @@ def test_object_rendering(
 
     renderer = treescope.active_renderer.get()
     # Render it to IR.
-    rendering = rendering_parts.build_full_line_with_annotations(
-        renderer.to_foldable_representation(
-            target, ignore_exceptions=ignore_exceptions
-        )
-    )
+    with warnings.catch_warnings():
+      if ignore_exceptions:
+        # Also ignore warnings due to ignoring exceptions to avoid cluttering
+        # logs.
+        warnings.simplefilter("ignore")
+      rendering = rendering_parts.build_full_line_with_annotations(
+          renderer.to_foldable_representation(
+              target, ignore_exceptions=ignore_exceptions
+          )
+      )
 
     # Collapse all foldables.
     layout_algorithms.expand_to_depth(rendering, 0)
@@ -676,8 +682,8 @@ def test_fallback_repr_pytree_node(self):
               <custom repr for UnknownPytreeNode: x=1234, y=5678>
                 #╭┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄╮
                 # PyTree children:
-                  GetAttrKey(name='x'): 1234,
-                  'custom_key': 5678,
+                  FrozenDataclassKey(name='x'): 1234,
+                  'string_key': 5678,
                 #╰┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄╯
               ,  # {object.__repr__(target[0])}
             ]"""),
@@ -712,9 +718,19 @@ def test_fallback_repr_after_error(self):
     ):
       renderer.to_foldable_representation(target)
 
-    rendering = rendering_parts.build_full_line_with_annotations(
-        renderer.to_foldable_representation(target, ignore_exceptions=True)
-    )
+    with self.assertWarnsRegex(
+        UserWarning,
+        "(.|\n)*".join([
+            re.escape("Ignoring error while formatting value of type"),
+            re.escape("ObjectWithCustomHandlerThatThrows"),
+            re.escape(
+                'raise RuntimeError("Simulated treescope_repr failure!")'
+            ),
+        ]),
+    ):
+      rendering = rendering_parts.build_full_line_with_annotations(
+          renderer.to_foldable_representation(target, ignore_exceptions=True)
+      )
 
     layout_algorithms.expand_to_depth(rendering, 0)
     self.assertEqual(
@@ -842,9 +858,17 @@ def test_failsafe_for_throw_in_repr(self):
     ):
       renderer.to_foldable_representation(target)
 
-    rendering = rendering_parts.build_full_line_with_annotations(
-        renderer.to_foldable_representation(target, ignore_exceptions=True)
-    )
+    with self.assertWarnsRegex(
+        UserWarning,
+        "(.|\n)*".join([
+            re.escape("Ignoring error while formatting value of type"),
+            re.escape("ObjectWithReprThatThrows"),
+            re.escape('raise RuntimeError("Simulated repr failure!")'),
+        ]),
+    ):
+      rendering = rendering_parts.build_full_line_with_annotations(
+          renderer.to_foldable_representation(target, ignore_exceptions=True)
+      )
 
     layout_algorithms.expand_to_depth(rendering, 0)
     self.assertEqual(