_realloc_island

PiperOrigin-RevId: 941618362
Change-Id: Idef1deb0ff4976de42aa585993988ae3363cd1fc

python/mujoco/bindings_test.py

@@ -699,6 +699,49 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     self.assertEmpty(self.data.contact)
     self.assertEmpty(self.data.efc_id)
 
+  def test_realloc_island(self):
+    # Test allocation on fresh data (on its own)
+    nisland = 2
+    nidof = 4
+    mujoco._functions._realloc_island(self.data, nisland=nisland, nidof=nidof)
+    self.assertEqual(self.data.nisland, nisland)
+    self.assertEqual(self.data.nidof, nidof)
+    self.assertEqual(self.data.island_nv.shape, (nisland,))
+    self.assertEqual(self.data.ifrc_smooth.shape, (nidof,))
+
+    # Test allocation after _realloc_con_efc
+    nefc = 10
+    mujoco._functions._realloc_con_efc(self.data, ncon=0, nefc=nefc)
+
+    nisland = 3
+    nidof = 5
+    mujoco._functions._realloc_island(self.data, nisland=nisland, nidof=nidof)
+
+    self.assertEqual(self.data.nisland, nisland)
+    self.assertEqual(self.data.nidof, nidof)
+    self.assertEqual(self.data.island_nv.shape, (nisland,))
+    self.assertEqual(self.data.ifrc_smooth.shape, (nidof,))
+
+    # Test re-allocation (calling it again with different sizes)
+    nisland2 = 4
+    nidof2 = 6
+    mujoco._functions._realloc_island(self.data, nisland=nisland2, nidof=nidof2)
+
+    self.assertEqual(self.data.nisland, nisland2)
+    self.assertEqual(self.data.nidof, nidof2)
+    self.assertEqual(self.data.island_nv.shape, (nisland2,))
+    self.assertEqual(self.data.ifrc_smooth.shape, (nidof2,))
+
+    # Test insufficient memory handling
+    expected_error = (
+        r'Insufficient arena memory, currently allocated memory=' +
+        r'"[0-9]+[A-Z]?". Increase using <size memory="X"/>.'
+    )
+    with self.assertRaisesRegex(mujoco.FatalError, expected_error):
+      mujoco._functions._realloc_island(self.data, 100000000, 100000000)
+    self.assertEqual(self.data.nisland, 0)
+    self.assertEqual(self.data.nidof, 0)
+
   def test_mj_struct_list_equality(self):
     model2 = mujoco.MjModel.from_xml_string(TEST_XML)
     data2 = mujoco.MjData(model2)