from __future__ import annotations
import os
import tempfile
import unittest
from unittest import mock
from rlm.bash import bash
from rlm import sleep_guard
class SleepGuardTests(unittest.TestCase):
def setUp(self) -> None:
sleep_guard.install_sleep_guard()
os.environ.pop("RLM_CELL_SLEEP_LIMIT", None)
def tearDown(self) -> None:
sleep_guard.install_sleep_guard()
os.environ.pop("RLM_CELL_SLEEP_LIMIT", None)
def test_long_literal_sleep_refused_with_options(self) -> None:
code = compile("import time; time.sleep(15)", "<cell-test-long>", "exec")
with self.assertRaises(ValueError) as ctx:
eval(code, {"__name__": "__main__"})
msg = str(ctx.exception)
self.assertIn("r = await h", msg)
self.assertIn("sentinel `bash()`", msg)
def test_short_sleep_allowed(self) -> None:
code = compile("import time; time.sleep(0.001)", "<cell-test-short>", "exec")
eval(code, {"__name__": "__main__"})
def test_library_frame_sleep_allowed(self) -> None:
with tempfile.TemporaryDirectory() as tmpdir:
mod_file = os.path.join(tmpdir, "helper_lib.py")
with open(mod_file, "w") as f:
f.write("import time\ndef run_sleep():\n time.sleep(15)\n")
import sys
sys.path.insert(0, tmpdir)
try:
import helper_lib
with mock.patch.object(sleep_guard, "_ORIG_TIME_SLEEP") as mock_orig:
code = compile("helper_lib.run_sleep()", "<cell-test-lib>", "exec")
eval(code, {"helper_lib": helper_lib, "__name__": "__main__"})
mock_orig.assert_called_once_with(15)
finally:
sys.path.remove(tmpdir)
sys.modules.pop("helper_lib", None)
def test_escape_hatch_allows_sleep(self) -> None:
os.environ["RLM_CELL_SLEEP_LIMIT"] = "0"
with mock.patch.object(sleep_guard, "_ORIG_TIME_SLEEP") as mock_orig:
code = compile("import time; time.sleep(15)", "<cell-test-escape>", "exec")
eval(code, {"__name__": "__main__"})
mock_orig.assert_called_once_with(15)
def test_cell_awaiting_handle_does_not_fire_guard(self) -> None:
import asyncio
async def run_cmd():
h = bash("echo test")
return await h
res = asyncio.run(run_cmd())
self.assertEqual(res.exit_code, 0)
if __name__ == "__main__":
unittest.main()