Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

tests/test_viser_update_policy.py

@@ -0,0 +1,116 @@
+"""Tests for Viser viewer update-policy helpers."""
+
+from dataclasses import dataclass
+from typing import Any
+from unittest.mock import MagicMock
+
+import mujoco
+from mjviser import ViserMujocoScene
+
+from mjlab.viewer.viser.overlays import ViserContactOverlays, ViserDebugOverlays
+from mjlab.viewer.viser.viewer import ViserPlayViewer
+
+
+@dataclass
+class _DummyScene:
+  env_idx: int
+  debug_visualization_enabled: bool
+  show_contact_points: bool = False
+  show_contact_forces: bool = False
+  needs_update: bool = False
+  clear_count: int = 0
+  clear_debug_count: int = 0
+
+  def clear(self) -> None:
+    self.clear_count += 1
+
+  def clear_debug_all(self) -> None:
+    self.clear_debug_count += 1
+
+
+class _DummyEnv:
+  def __init__(self, unwrapped: Any):
+    self._unwrapped = unwrapped
+
+  @property
+  def unwrapped(self) -> Any:
+    return self._unwrapped
+
+
+def test_should_update_cameras():
+  assert ViserPlayViewer._should_update_cameras(paused=False, has_pending_updates=False)
+  assert not ViserPlayViewer._should_update_cameras(
+    paused=True, has_pending_updates=False
+  )
+  assert ViserPlayViewer._should_update_cameras(paused=True, has_pending_updates=True)
+
+
+def test_should_submit_scene_update():
+  # Odd ticks are skipped to keep scene submits around 30Hz.
+  assert not ViserPlayViewer._should_submit_scene_update(
+    counter=1, paused=False, has_pending_updates=True
+  )
+  # Running: submit on even ticks regardless of pending flags.
+  assert ViserPlayViewer._should_submit_scene_update(
+    counter=2, paused=False, has_pending_updates=False
+  )
+  # Paused: submit only with pending updates.
+  assert not ViserPlayViewer._should_submit_scene_update(
+    counter=2, paused=True, has_pending_updates=False
+  )
+  assert ViserPlayViewer._should_submit_scene_update(
+    counter=2, paused=True, has_pending_updates=True
+  )
+
+
+def test_scene_requires_live_refresh():
+  """Base scene sets needs_update when contact overlays are visible."""
+  model = mujoco.MjModel.from_xml_string("<mujoco><worldbody/></mujoco>")
+  server = MagicMock()
+  scene = ViserMujocoScene(server, model, num_envs=1)
+
+  scene.show_contact_points = False
+  scene.show_contact_forces = False
+  scene.needs_update = False
+  scene._last_body_xpos = None  # No cached data, refresh is a no-op.
+  scene.refresh_visualization()
+  assert not scene.needs_update
+
+  scene.show_contact_points = True
+  assert scene.show_contact_points or scene.show_contact_forces
+
+
+def test_debug_overlays_env_switch_and_queue():
+  """env_switch respects enabled flag; queue clears + calls update_visualizers."""
+  unwrapped = MagicMock(spec=["update_visualizers"])
+  env = _DummyEnv(unwrapped)
+  scene = _DummyScene(env_idx=0, debug_visualization_enabled=False)
+  overlays = ViserDebugOverlays(env=env, scene=scene)
+
+  # Disabled: env switch does not clear.
+  overlays.on_env_switch()
+  assert scene.clear_debug_count == 0
+
+  # Enabled: env switch clears, queue dispatches.
+  scene.debug_visualization_enabled = True
+  overlays.on_env_switch()
+  assert scene.clear_debug_count == 1
+
+  overlays.queue()
+  assert scene.clear_count == 1
+  unwrapped.update_visualizers.assert_called_once_with(scene)
+
+
+def test_contact_overlays_env_switch():
+  """env_switch requests scene update only when contacts are enabled."""
+  scene = _DummyScene(env_idx=0, debug_visualization_enabled=False)
+  overlays = ViserContactOverlays(scene=scene)
+
+  assert not overlays.is_enabled()
+  overlays.on_env_switch()
+  assert not scene.needs_update
+
+  scene.show_contact_points = True
+  assert overlays.is_enabled()
+  overlays.on_env_switch()
+  assert scene.needs_update