keep render util test setup

Restore the existing Warp/CUDA-gated render_util test setup so the review change stays minimal while preserving the new batched shape coverage.

Made-with: Cursor

mjx/mujoco/mjx/_src/render_util_test.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import contextlib
+import os
 from unittest import mock
 
 from absl.testing import absltest
@@ -20,9 +20,13 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 
+from mujoco.mjx._src import io
 from mujoco.mjx._src import render_util
+import mujoco.mjx.warp as mjxw
 from mujoco.mjx.warp.render_context import RenderContextPytree
 
+_FORCE_TEST = os.environ.get('MJX_WARP_FORCE_TEST', '0') == '1'
+
 
 def _fake_render_context(ncam, width, height):
   """Fake RenderContext for testing."""
@@ -36,25 +40,26 @@ def _fake_render_context(ncam, width, height):
   return rc
 
 
-@contextlib.contextmanager
-def _mock_render_runtime(warp_rc):
-  with mock.patch.object(render_util.mjxw, 'WARP_INSTALLED', True):
-    with mock.patch.dict(
-        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
-        {(0, None): warp_rc},
-    ):
-      yield
-
-
 class RenderUtilTest(absltest.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    if not _FORCE_TEST:
+      if not mjxw.WARP_INSTALLED:
+        self.skipTest('Warp not installed.')
+      if not io.has_cuda_gpu_device():
+        self.skipTest('No CUDA GPU device available.')
+
   def test_get_rgb(self):
     width, height = 4, 4
     warp_rc = _fake_render_context(1, width, height)
     rc = mock.MagicMock(spec=RenderContextPytree, key=0)
     rgb_data = jnp.zeros((width * height,), dtype=jnp.uint32)
 
-    with _mock_render_runtime(warp_rc):
+    with mock.patch.dict(
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
+        {(0, None): warp_rc},
+    ):
       rgb = jax.jit(render_util.get_rgb, static_argnums=(0, 1))(rc, 0, rgb_data)
 
     self.assertEqual(rgb.shape, (height, width, 3))
@@ -64,7 +69,10 @@ class RenderUtilTest(absltest.TestCase):
     warp_rc = _fake_render_context(1, width, height)
     rc = mock.MagicMock(spec=RenderContextPytree, key=0)
 
-    with _mock_render_runtime(warp_rc):
+    with mock.patch.dict(
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
+        {(0, None): warp_rc},
+    ):
       for leading_shape in ((1,), (3,), (2, 3)):
         with self.subTest(leading_shape=leading_shape):
           rgb_data = jnp.zeros(
@@ -82,7 +90,10 @@ class RenderUtilTest(absltest.TestCase):
     rc = mock.MagicMock(spec=RenderContextPytree, key=0)
     rgb_data = jnp.zeros((nworld, width * height), dtype=jnp.uint32)
 
-    with _mock_render_runtime(warp_rc):
+    with mock.patch.dict(
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
+        {(0, None): warp_rc},
+    ):
       rgb = jax.jit(
           jax.vmap(render_util.get_rgb, in_axes=(None, None, 0)),
           static_argnums=(0, 1),
@@ -96,7 +107,10 @@ class RenderUtilTest(absltest.TestCase):
     rc = mock.MagicMock(spec=RenderContextPytree, key=0)
     depth_data = jnp.zeros((width * height,), dtype=jnp.float32)
 
-    with _mock_render_runtime(warp_rc):
+    with mock.patch.dict(
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
+        {(0, None): warp_rc},
+    ):
       depth = jax.jit(render_util.get_depth, static_argnums=(0, 1, 3))(
           rc, 0, depth_data, 5.0
       )
@@ -108,7 +122,10 @@ class RenderUtilTest(absltest.TestCase):
     warp_rc = _fake_render_context(1, width, height)
     rc = mock.MagicMock(spec=RenderContextPytree, key=0)
 
-    with _mock_render_runtime(warp_rc):
+    with mock.patch.dict(
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
+        {(0, None): warp_rc},
+    ):
       for leading_shape in ((1,), (3,), (2, 3)):
         with self.subTest(leading_shape=leading_shape):
           depth_data = jnp.zeros(
@@ -126,7 +143,10 @@ class RenderUtilTest(absltest.TestCase):
     rc = mock.MagicMock(spec=RenderContextPytree, key=0)
     depth_data = jnp.zeros((nworld, width * height), dtype=jnp.float32)
 
-    with _mock_render_runtime(warp_rc):
+    with mock.patch.dict(
+        'mujoco.mjx.warp.render_context._MJX_RENDER_CONTEXT_BUFFERS',
+        {(0, None): warp_rc},
+    ):
       depth = jax.jit(
           jax.vmap(render_util.get_depth, in_axes=(None, None, 0, None)),
           static_argnums=(0, 1, 3),