diff --git a/rock/sdk/sandbox/client.py b/rock/sdk/sandbox/client.py index f19c9b0fc..88ea07ea2 100644 --- a/rock/sdk/sandbox/client.py +++ b/rock/sdk/sandbox/client.py @@ -882,8 +882,22 @@ def _generate_utc_iso_time(self): return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") async def close_session(self, request: CloseSessionRequest) -> CloseSessionResponse: - # TODO: implement this - pass + url = f"{self._url}/close_session" + headers = self._build_headers() + data = { + "sandbox_id": self.sandbox_id, + **request.model_dump(), + } + try: + response = await HttpUtils.post(url, headers, data) + except Exception as e: + raise Exception(f"Failed to close session: {str(e)}, post url {url}") + + logging.debug(f"Close session response: {response}") + if "Success" != response.get("status"): + raise Exception(f"Failed to close session: {response}") + result: dict = response.get("result") # type: ignore + return CloseSessionResponse(**result) async def close(self) -> CloseResponse: await self.stop() diff --git a/tests/integration/sdk/sandbox/test_close_session.py b/tests/integration/sdk/sandbox/test_close_session.py new file mode 100644 index 000000000..ac5d54dab --- /dev/null +++ b/tests/integration/sdk/sandbox/test_close_session.py @@ -0,0 +1,26 @@ +import pytest +import os +from pathlib import Path +from rock.sdk.sandbox.client import Sandbox +from rock.actions import CreateBashSessionRequest, CloseBashSessionRequest, BashAction +from tests.integration.conftest import SKIP_IF_NO_DOCKER + +# Set writable status directory for sandbox deployment +os.environ["ROCK_SERVICE_STATUS_DIR"] = "/tmp/rock_status" + + +@pytest.mark.need_admin +@SKIP_IF_NO_DOCKER +@pytest.mark.asyncio +async def test_sdk_close_session(sandbox_instance: Sandbox): + session_name = "test-close-session" + await sandbox_instance.create_session(CreateBashSessionRequest(session=session_name, session_type="bash")) + + obs = await sandbox_instance.arun(cmd="echo 'alive'", session=session_name) + assert "alive" in obs.output + + resp = await sandbox_instance.close_session(CloseBashSessionRequest(session=session_name, session_type="bash")) + assert resp is not None + + with pytest.raises(Exception): + await sandbox_instance.arun(cmd="echo 'should fail'", session=session_name)