Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
4.9 kB · 142 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143from __future__ import annotations
import asyncioimport osimport sysimport unittest
SRC = os.path.join(os.path.dirname(__file__), "..", "src")if SRC not in sys.path: sys.path.insert(0, SRC)
import remotefrom 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()