from __future__ import annotations
import asyncio
import os
import sys
import unittest
SRC = os.path.join(os.path.dirname(__file__), "..", "src")
if SRC not in sys.path:
sys.path.insert(0, SRC)
import remote
from remote import RemoteExecutionError, RemoteHandle, RemoteKernelHandle
class RemoteKernelHandleTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
env = {
**os.environ,
"PYTHONPATH": SRC + os.pathsep + os.environ.get("PYTHONPATH", ""),
}
self.proc = await asyncio.create_subprocess_exec(
sys.executable,
"-m",
"rlm.repl",
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env,
)
assert self.proc.stdout is not None
ready_line = await self.proc.stdout.readline()
self.assertTrue(ready_line.startswith(b'{"event":"ready"'))
self.handle = RemoteKernelHandle("localhost", self.proc, sys.executable, "/tmp")
async def asyncTearDown(self):
await self.handle.close()
async def test_imports_and_call_lambda(self):
await self.handle.imports("math")
res = await self.handle.call(lambda: math.sqrt(144))
self.assertEqual(res, 12.0)
async def test_chained_imports(self):
res = await (await self.handle.imports("platform")).call(lambda: bool(platform.system()))
self.assertTrue(res)
async def test_repeated_imports_and_reload(self):
await self.handle.imports("math")
self.assertIn("math", self.handle._imported_modules)
# Second import is a client-side no-op
await self.handle.imports("math")
res = await self.handle.call(lambda: math.pi > 3.0)
self.assertTrue(res)
# Explicit reload forces re-execution
await self.handle.imports("math", reload=True)
res2 = await self.handle.call(lambda: math.pi > 3.0)
self.assertTrue(res2)
async def test_call_with_args_and_kwargs(self):
def multiply_and_add(x, y, offset=0):
return x * y + offset
res = await self.handle.call(multiply_and_add, 6, 7, offset=5)
self.assertEqual(res, 47)
async def test_call_async_function(self):
async def async_sum(a, b):
await asyncio.sleep(0.01)
return a + b
res = await self.handle.call(async_sum, 15, 27)
self.assertEqual(res, 42)
async def test_call_exception_handling(self):
with self.assertRaises(RemoteExecutionError) as ctx:
await self.handle.call(lambda: 1 / 0)
self.assertIn("ZeroDivisionError", str(ctx.exception))
async def test_handle_liveness_and_poll(self):
async def slow_work():
await asyncio.sleep(0.05)
return "done!"
h = self.handle.call(slow_work)
self.assertIsInstance(h, RemoteHandle)
self.assertTrue(h.running)
self.assertIsNone(h.poll())
res = await h
self.assertEqual(res, "done!")
self.assertFalse(h.running)
self.assertEqual(h.poll(), "done!")
async def test_handle_output_streaming(self):
def print_and_return():
print("step 1: started")
print("step 2: finished")
return 123
h = self.handle.call(print_and_return)
res = await h
self.assertEqual(res, 123)
self.assertIn("step 1: started", h.text())
self.assertIn("step 2: finished", h.text())
self.assertIn("step 1: started", h.text())
self.assertIn("step 2", h.tail(1))
async def test_first_class_remote_bash_handle(self):
h = self.handle.bash("echo 'first class remote bash'")
self.assertIsInstance(h, RemoteHandle)
res = await h
self.assertEqual(res.exit_code, 0)
self.assertIn("first class remote bash", res.text())
async def test_first_class_remote_bash_handle_with_input(self):
h = self.handle.bash("cat", input="remote stdin string")
self.assertIsInstance(h, RemoteHandle)
res = await h
self.assertEqual(res.exit_code, 0)
self.assertEqual(res.text(), "remote stdin string")
async def test_first_class_remote_bash_handle_with_bash_handle_input(self):
from rlm.bash import bash as local_bash
local_h = await local_bash("echo 'piped from local handle'")
h = self.handle.bash("cat", input=local_h)
res = await h
self.assertEqual(res.exit_code, 0)
self.assertIn("piped from local handle", res.text())
async def test_async_context_manager(self):
async with self.handle as rem:
await rem.imports("sys")
res = await rem.call(lambda: sys.platform)
self.assertIsInstance(res, str)
self.assertTrue(self.handle.closed)
if __name__ == "__main__":
unittest.main()