def test_sh_producing_large_outputs(self):
"""Process a few megabytes of output and assert that:
- It will consume not more than 10 megabytes (assuming also output capturing in tests by io.capture_descriptors())
- The whole output would be printed correctly
"""
self.maxDiff = None # unittest setting
task = InitTask()
task._io = IO()
io = IO()
out = StringIO()
text = "History isn't made by kings and politicians, it is made by us."
memory_before = psutil.Process(
os.getpid()).memory_info().rss / 1024 / 1024
with io.capture_descriptors(stream=out, enable_standard_out=False):
task.py('''
for i in range(0, 1024 * 128):
print("''' + text + '''")
''')
iterations = 1024 * 128
text_with_newlines_length = len(text) + 2 # \r + \n
memory_after = psutil.Process(
os.getpid()).memory_info().rss / 1024 / 1024
self.assertEqual(iterations * text_with_newlines_length,
len(out.getvalue()))
self.assertLessEqual(
memory_after - memory_before,
16,
msg='Expected less than 16 megabytes of memory usage')
def test_py_uses_sudo_when_become_specified(self):
"""Expect that sudo with proper parameters is used"""
task = InitTask()
task._io = IO()
with unittest.mock.patch(
'rkd.taskutil.check_output') as mocked_subprocess:
mocked_subprocess: unittest.mock.MagicMock
task.py(code='print("test")', capture=True, become='root')
self.assertIn('sudo -E -u root python',
mocked_subprocess.call_args[0][0])
def test_py_inherits_environment_variables(self):
os.putenv('PY_INHERITS_ENVIRONMENT_VARIABLES', 'should')
task = InitTask()
task._io = IO()
out = task.py(
code=
'import os; print("ENV VALUE IS: " + str(os.environ["PY_INHERITS_ENVIRONMENT_VARIABLES"]))',
capture=True)
self.assertEqual('ENV VALUE IS: should\n', out)
def test_py_executes_python_scripts_without_specifying_script_path(self):
"""Simply - check basic successful case - executing a Python code"""
task = InitTask()
task._io = IO()
out = task.py('''
import os
print(os)
''', capture=True)
self.assertIn("<module 'os' from", out)
def test_py_executes_a_custom_python_script(self):
"""Check that script from specified file in 'script_path' parameter will be executed
And the code will be passed to that script as stdin.
"""
task = InitTask()
task._io = IO()
with NamedTemporaryFile() as temp_file:
temp_file.write(b'import sys; print(sys.argv[1])')
temp_file.flush()
out = task.py('',
capture=True,
script_path=temp_file.name,
arguments='Hello!')
self.assertEqual('Hello!\n', out)