From 80cd4caa1b670d35db43820174f59722949ad390 Mon Sep 17 00:00:00 2001 From: Grain Team Date: Tue, 12 May 2026 14:00:41 -0700 Subject: [PATCH] internal PiperOrigin-RevId: 914465125 --- grain/_src/core/multiprocessing_test.py | 50 +++++++++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 grain/_src/core/multiprocessing_test.py diff --git a/grain/_src/core/multiprocessing_test.py b/grain/_src/core/multiprocessing_test.py new file mode 100644 index 000000000..b351f5570 --- /dev/null +++ b/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()