diff --git a/setup.py b/setup.py index 2cc1d9d..36049ca 100644 --- a/setup.py +++ b/setup.py @@ -62,6 +62,7 @@ 'privileged = rocker.extensions:Privileged', 'pulse = rocker.extensions:PulseAudio', 'rmw = rocker.rmw_extension:RMW', + 'shm_size = rocker.extensions:ShmSize', 'ssh = rocker.ssh_extension:Ssh', 'ulimit = rocker.ulimit_extension:Ulimit', 'user = rocker.extensions:User', diff --git a/src/rocker/extensions.py b/src/rocker/extensions.py index 3e22b71..bee9363 100644 --- a/src/rocker/extensions.py +++ b/src/rocker/extensions.py @@ -468,3 +468,27 @@ def register_arguments(parser, defaults): default=defaults.get(GroupAdd.get_name(), None), action='append', help="Add additional groups to join.") + +class ShmSize(RockerExtension): + @staticmethod + def get_name(): + return 'shm_size' + + def __init__(self): + self.name = ShmSize.get_name() + + def get_preamble(self, cliargs): + return '' + + def get_docker_args(self, cliargs): + args = '' + shm_size = cliargs.get('shm_size', None) + if shm_size: + args += f' --shm-size {shm_size} ' + return args + + @staticmethod + def register_arguments(parser, defaults={}): + parser.add_argument('--shm-size', + default=defaults.get('shm_size', None), + help="Set the size of the shared memory for the container (e.g., 512m, 1g).") \ No newline at end of file diff --git a/test/test_extension.py b/test/test_extension.py index d264c35..1ad88d8 100644 --- a/test/test_extension.py +++ b/test/test_extension.py @@ -617,3 +617,35 @@ def test_group_add_extension(self): args = p.get_docker_args(mock_cliargs) self.assertIn('--group-add sudo', args) self.assertIn('--group-add docker', args) + +class ShmSizeExtensionTest(unittest.TestCase): + + def setUp(self): + # Work around interference between empy Interpreter + # stdout proxy and test runner. empy installs a proxy on stdout + # to be able to capture the information. + # And the test runner creates a new stdout object for each test. + # This breaks empy as it assumes that the proxy has persistent + # between instances of the Interpreter class + # empy will error with the exception + # "em.Error: interpreter stdout proxy lost" + em.Interpreter._wasProxyInstalled = False + + @pytest.mark.docker + def test_shm_size_extension(self): + plugins = list_plugins() + shm_size_plugin = plugins['shm_size'] + self.assertEqual(shm_size_plugin.get_name(), 'shm_size') + + p = shm_size_plugin() + self.assertTrue(plugin_load_parser_correctly(shm_size_plugin)) + + mock_cliargs = {} + self.assertEqual(p.get_snippet(mock_cliargs), '') + self.assertEqual(p.get_preamble(mock_cliargs), '') + args = p.get_docker_args(mock_cliargs) + self.assertNotIn('--shm-size', args) + + mock_cliargs = {'shm_size': '12g'} + args = p.get_docker_args(mock_cliargs) + self.assertIn('--shm-size 12g', args) \ No newline at end of file