internal

grain/_src/core/multiprocessing_test.py

@@ -0,0 +1,50 @@
+import importlib
+import os
+from unittest import mock
+
+from absl.testing import absltest
+import multiprocessing as grain_mp
+
+# Use importlib to load shared_memory dynamically to bypass false-positive
+# strict deps failures regarding missing direct dependencies for standard
+# library.
+shared_memory = importlib.import_module("multiprocessing.shared_memory")
+
+
+class MultiprocessingTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.orig_shm_name = grain_mp._endoscope_shm_name
+
+  def tearDown(self):
+    grain_mp._endoscope_shm_name = self.orig_shm_name
+    super().tearDown()
+
+  def test_get_shm_name_is_random(self):
+    name1 = grain_mp._get_shm_name()
+    grain_mp._endoscope_shm_name = None
+    name2 = grain_mp._get_shm_name()
+    self.assertNotEqual(name1, name2)
+    self.assertIn(str(os.getpid()), name1)
+
+  def test_initialize_endoscope_shm_size_mismatch(self):
+    test_name = f"test_grain_shm_{os.getpid()}"
+    grain_mp._endoscope_shm_name = test_name
+
+    mock_shm = mock.create_autospec(shared_memory.SharedMemory, instance=True)
+    mock_shm.size = 4  # Small size
+
+    with mock.patch.object(
+        shared_memory, "SharedMemory", autospec=True
+    ) as mock_sm:
+      mock_sm.return_value = mock_shm
+
+      with self.assertRaisesRegex(
+          ValueError, "Shared memory segment size mismatch: 4 != "
+      ):
+        grain_mp._initialize_endoscope_shm()
+
+
+if __name__ == "__main__":
+  grain_absltest.main()