Fix flaky viz test in test_examples_viz.py (mesa#3039)

vedantvakharia · web-flow · commit e9054adbb9ef · 2025-12-29T17:51:55.000+01:00
diff --git a/tests/examples/test_examples_viz.py b/tests/examples/test_examples_viz.py
@@ -72,12 +72,12 @@ def run_model_test(
         # Display and capture the updated visualizations
         display(space_viz)
         page_session.wait_for_selector("img")
-        changed_space = page_session.locator("img").screenshot()
+        changed_space = page_session.locator("img").first.screenshot()
 
         if measure_config:
             display(graph_viz)
             page_session.wait_for_selector("img")
-            changed_graph = page_session.locator("img").screenshot()
+            changed_graph = page_session.locator("img").last.screenshot()
 
         # Convert screenshots to base64 for comparison
         initial_space_encoding = base64.b64encode(initial_space).decode()
