mirror of
https://github.com/NVIDIA/OpenShell.git
synced 2026-10-02 07:34:45 +08:00
Matthew Grossman
76cfd0e31d
* test(python): synchronize interactive exec TTY readiness Wait for the complete readiness marker before streaming stdin so PTY echo cannot split the separately written TTY flags. Preserve pipe stream separation and verify consumed stdin and both output sentinels in TTY mode. Fixes #4075 Signed-off-by: Matthew Grossman <mgrossman@nvidia.com> * test(python): reuse interactive exec readiness marker Signed-off-by: Matthew Grossman <mgrossman@nvidia.com> --------- Signed-off-by: Matthew Grossman <mgrossman@nvidia.com>
386 lines
14 KiB
Python
386 lines
14 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import threading
|
|
import uuid
|
|
from typing import TYPE_CHECKING
|
|
|
|
from google.protobuf import duration_pb2
|
|
|
|
from openshell._proto import datamodel_pb2, openshell_pb2, sandbox_pb2
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
from openshell import Sandbox, SandboxClient, WorkspaceClient
|
|
|
|
|
|
def test_mutation_replay_preserves_sandbox_lifecycle_and_replacement(
|
|
sandbox_client: SandboxClient,
|
|
) -> None:
|
|
name = f"replay-{uuid.uuid4().hex[:8]}"
|
|
scope = "default"
|
|
stub = sandbox_client._stub
|
|
create = openshell_pb2.CreateSandboxRequest(
|
|
name=name,
|
|
spec=openshell_pb2.SandboxSpec(),
|
|
workspace_scope=datamodel_pb2.WorkspaceSelector(workspace=scope),
|
|
request_id=str(uuid.uuid4()),
|
|
)
|
|
|
|
def replay(method, request):
|
|
result, call = method.with_call(request, timeout=60)
|
|
assert dict(call.initial_metadata())["openshell-replayed"] == "true"
|
|
return result
|
|
|
|
try:
|
|
original = stub.CreateSandbox(create, timeout=60).sandbox.metadata.id
|
|
sandbox_client.wait_ready(name, workspace="default", timeout_seconds=300)
|
|
assert replay(stub.CreateSandbox, create).sandbox.metadata.id == original
|
|
stop = openshell_pb2.StopSandboxRequest(
|
|
name=name,
|
|
workspace_scope=datamodel_pb2.WorkspaceSelector(workspace=scope),
|
|
request_id=str(uuid.uuid4()),
|
|
)
|
|
stub.StopSandbox(stop, timeout=60)
|
|
sandbox_client.wait_stopped(name, workspace="default", timeout_seconds=120)
|
|
assert replay(stub.StopSandbox, stop).sandbox.metadata.id == original
|
|
start = openshell_pb2.StartSandboxRequest(
|
|
name=name,
|
|
workspace_scope=datamodel_pb2.WorkspaceSelector(workspace=scope),
|
|
request_id=str(uuid.uuid4()),
|
|
)
|
|
stub.StartSandbox(start, timeout=60)
|
|
sandbox_client.wait_ready(name, workspace="default", timeout_seconds=300)
|
|
assert replay(stub.StartSandbox, start).sandbox.metadata.id == original
|
|
update = openshell_pb2.UpdateConfigRequest(
|
|
sandbox=name,
|
|
workspace_scope=datamodel_pb2.WorkspaceSelector(workspace=scope),
|
|
setting_key="ocsf_json_enabled",
|
|
setting_value=sandbox_pb2.SettingValue(bool_value=True),
|
|
request_id=str(uuid.uuid4()),
|
|
)
|
|
updated = stub.UpdateConfig(update, timeout=30)
|
|
assert replay(stub.UpdateConfig, update) == updated
|
|
delete = openshell_pb2.DeleteSandboxRequest(
|
|
name=name,
|
|
workspace_scope=datamodel_pb2.WorkspaceSelector(workspace=scope),
|
|
request_id=str(uuid.uuid4()),
|
|
)
|
|
deleted = stub.DeleteSandbox(delete, timeout=60)
|
|
assert deleted.sandbox_id == original
|
|
sandbox_client.wait_deleted(name, workspace="default", timeout_seconds=120)
|
|
replacement = sandbox_client.create(workspace="default", name=name)
|
|
assert replacement.id != original
|
|
assert replay(stub.DeleteSandbox, delete) == deleted
|
|
assert sandbox_client.get(name, workspace="default").id == replacement.id
|
|
finally:
|
|
with contextlib.suppress(Exception):
|
|
sandbox_client.delete(name, workspace="default", allow_missing=True)
|
|
sandbox_client.wait_deleted(name, workspace="default", timeout_seconds=120)
|
|
|
|
|
|
def test_sandbox_api_crud_and_exec(
|
|
sandbox: Callable[..., Sandbox],
|
|
sandbox_client: SandboxClient,
|
|
) -> None:
|
|
class _FileOps:
|
|
def write(self, path: str, content: str) -> None:
|
|
from pathlib import Path
|
|
|
|
Path(path).write_text(content)
|
|
|
|
def read(self, path: str) -> str:
|
|
from pathlib import Path
|
|
|
|
return Path(path).read_text()
|
|
|
|
with sandbox(delete_on_exit=True) as sb:
|
|
assert sb.id
|
|
# Server auto-generates a petname (e.g. "feasible-retriever")
|
|
assert sb.sandbox.name
|
|
parts = sb.sandbox.name.split("-")
|
|
assert len(parts) == 2, (
|
|
f"expected petname with 2 parts, got {sb.sandbox.name!r}"
|
|
)
|
|
assert all(p.isalpha() and p.islower() for p in parts)
|
|
|
|
fetched = sandbox_client.get(sb.sandbox.name, workspace="default")
|
|
assert fetched.id == sb.id
|
|
|
|
ids = set(sandbox_client.list_ids(workspace="default", page_size=100))
|
|
assert sb.id in ids
|
|
|
|
result = sb.exec(["python", "-c", "print('sandbox-ok')"])
|
|
assert result.exit_code == 0
|
|
assert "sandbox-ok" in result.stdout
|
|
|
|
file_ops = _FileOps()
|
|
create_file = sb.exec_python(
|
|
file_ops.write,
|
|
args=("/sandbox/exec-persistence.txt", "ok"),
|
|
)
|
|
assert create_file.exit_code == 0
|
|
|
|
verify_file = sb.exec_python(
|
|
file_ops.read, args=("/sandbox/exec-persistence.txt",)
|
|
)
|
|
assert verify_file.exit_code == 0
|
|
assert verify_file.stdout.strip() == "ok"
|
|
|
|
|
|
def test_sandbox_interactive_exec_honors_tty(
|
|
sandbox: Callable[..., Sandbox],
|
|
sandbox_client: SandboxClient,
|
|
) -> None:
|
|
ready_marker = b"tty-ready\n"
|
|
stdin_sentinel = b"streamed-stdin-sentinel"
|
|
stdout_sentinel = b"stdout-sentinel"
|
|
stderr_sentinel = b"stderr-sentinel"
|
|
|
|
def exec_interactive(sandbox_name: str, *, tty: bool) -> tuple[bytes, bytes]:
|
|
request = openshell_pb2.ExecSandboxInput(
|
|
start=openshell_pb2.ExecSandboxRequest(
|
|
sandbox=sandbox_name,
|
|
workspace_scope=datamodel_pb2.WorkspaceSelector(workspace="default"),
|
|
command=[
|
|
"/bin/sh",
|
|
"-c",
|
|
"[ -t 0 ] && printf T || printf N; "
|
|
"[ -t 1 ] && printf T || printf N; "
|
|
"[ -t 2 ] && printf T || printf N; "
|
|
f"printf '\\n%s' '{ready_marker.decode()}'; "
|
|
"IFS= read -r stdin_value; "
|
|
"printf 'stdin:%s\\n' \"$stdin_value\"; "
|
|
"printf 'stdout-sentinel\\n'; "
|
|
"printf 'stderr-sentinel\\n' >&2",
|
|
],
|
|
tty=tty,
|
|
execution_timeout=duration_pb2.Duration(seconds=20),
|
|
)
|
|
)
|
|
|
|
ready = threading.Event()
|
|
done = threading.Event()
|
|
|
|
def requests():
|
|
yield request
|
|
# PTY input echo can split the separate TTY flag writes. Wait for
|
|
# the complete marker, including its newline, before sending input.
|
|
if not ready.wait(timeout=20) or done.is_set():
|
|
return
|
|
yield openshell_pb2.ExecSandboxInput(stdin=stdin_sentinel + b"\n")
|
|
done.wait(timeout=30)
|
|
|
|
stdout = bytearray()
|
|
stderr = bytearray()
|
|
exit_code: int | None = None
|
|
try:
|
|
events = sandbox_client._stub.ExecSandboxInteractive(requests(), timeout=30)
|
|
for event in events:
|
|
payload = event.WhichOneof("payload")
|
|
if payload == "stdout":
|
|
stdout.extend(event.stdout.data)
|
|
if ready_marker in stdout.replace(b"\r\n", b"\n"):
|
|
ready.set()
|
|
elif payload == "stderr":
|
|
stderr.extend(event.stderr.data)
|
|
elif payload == "exit":
|
|
exit_code = int(event.exit.exit_code)
|
|
finally:
|
|
done.set()
|
|
# Unblock a request iterator waiting for readiness on early exit
|
|
# or RPC failure without sending input after the call has ended.
|
|
ready.set()
|
|
|
|
assert exit_code == 0
|
|
return bytes(stdout), bytes(stderr)
|
|
|
|
with sandbox(delete_on_exit=True) as sb:
|
|
stdout, stderr = exec_interactive(sb.sandbox.name, tty=False)
|
|
assert b"NNN" in stdout.splitlines()
|
|
assert b"stdin:" + stdin_sentinel in stdout
|
|
assert stdout_sentinel in stdout
|
|
assert stdout_sentinel not in stderr
|
|
assert stderr_sentinel in stderr
|
|
assert stderr_sentinel not in stdout
|
|
|
|
stdout, stderr = exec_interactive(sb.sandbox.name, tty=True)
|
|
terminal_output = stdout + stderr
|
|
assert b"TTT" in terminal_output.splitlines()
|
|
assert b"stdin:" + stdin_sentinel in terminal_output
|
|
assert stdout_sentinel in terminal_output
|
|
assert stderr_sentinel in terminal_output
|
|
|
|
|
|
def test_interactive_exec_drains_output_after_request_eof(
|
|
sandbox: Callable[..., Sandbox],
|
|
sandbox_client: SandboxClient,
|
|
) -> None:
|
|
with sandbox(delete_on_exit=True) as sb:
|
|
|
|
def requests():
|
|
yield openshell_pb2.ExecSandboxInput(
|
|
start=openshell_pb2.ExecSandboxRequest(
|
|
sandbox=sb.sandbox.name,
|
|
workspace_scope=datamodel_pb2.WorkspaceSelector(workspace="default"),
|
|
command=[
|
|
"/bin/sh",
|
|
"-c",
|
|
"input=$(cat); sleep 0.2; "
|
|
"printf '%s' \"$input\"; printf 'drained-stderr' >&2; exit 7",
|
|
],
|
|
tty=False,
|
|
execution_timeout=duration_pb2.Duration(seconds=20),
|
|
)
|
|
)
|
|
yield openshell_pb2.ExecSandboxInput(stdin=b"drained-stdout")
|
|
# End requests before the command emits output. The receive
|
|
# direction must survive long enough to drain both output streams.
|
|
|
|
stdout = bytearray()
|
|
stderr = bytearray()
|
|
exit_codes: list[int] = []
|
|
for event in sandbox_client._stub.ExecSandboxInteractive(
|
|
requests(), timeout=30
|
|
):
|
|
assert not exit_codes, "received another event after terminal exit"
|
|
payload = event.WhichOneof("payload")
|
|
if payload == "stdout":
|
|
stdout.extend(event.stdout.data)
|
|
elif payload == "stderr":
|
|
stderr.extend(event.stderr.data)
|
|
elif payload == "exit":
|
|
exit_codes.append(event.exit.exit_code)
|
|
assert stdout == b"drained-stdout"
|
|
assert stderr == b"drained-stderr"
|
|
assert exit_codes == [7]
|
|
|
|
|
|
def test_list_scoped_and_for_all_workspaces(
|
|
sandbox_client: SandboxClient,
|
|
workspace_client: WorkspaceClient,
|
|
) -> None:
|
|
import contextlib
|
|
import uuid
|
|
|
|
suffix = uuid.uuid4().hex[:8]
|
|
other_ws = f"list-ws-{suffix}"
|
|
created_default: list[str] = []
|
|
created_other: list[str] = []
|
|
|
|
try:
|
|
workspace_client.create(other_ws)
|
|
|
|
ref_default = sandbox_client.create(
|
|
workspace="default", name=f"ls-def-{suffix}"
|
|
)
|
|
created_default.append(ref_default.name)
|
|
|
|
ref_other = sandbox_client.create(workspace=other_ws, name=f"ls-oth-{suffix}")
|
|
created_other.append(ref_other.name)
|
|
|
|
default_ids = set(sandbox_client.list_ids(workspace="default"))
|
|
assert ref_default.id in default_ids
|
|
assert ref_other.id not in default_ids
|
|
|
|
other_ids = set(sandbox_client.list_ids(workspace=other_ws))
|
|
assert ref_other.id in other_ids
|
|
assert ref_default.id not in other_ids
|
|
|
|
all_ids = set(sandbox_client.list_ids_for_all_workspaces())
|
|
assert ref_default.id in all_ids
|
|
assert ref_other.id in all_ids
|
|
finally:
|
|
for name in created_default:
|
|
with contextlib.suppress(Exception):
|
|
sandbox_client.delete(name, workspace="default")
|
|
sandbox_client.wait_deleted(name, workspace="default")
|
|
for name in created_other:
|
|
with contextlib.suppress(Exception):
|
|
sandbox_client.delete(name, workspace=other_ws)
|
|
sandbox_client.wait_deleted(name, workspace=other_ws)
|
|
with contextlib.suppress(Exception):
|
|
workspace_client.delete(other_ws)
|
|
|
|
|
|
def test_sandbox_labels_and_selectors(sandbox_client: SandboxClient) -> None:
|
|
import contextlib
|
|
import uuid
|
|
|
|
suffix = uuid.uuid4().hex[:8]
|
|
job_a = f"lbl-a-{suffix}"
|
|
job_b = f"lbl-b-{suffix}"
|
|
group_selector = f"aiq-test={suffix}"
|
|
primary_selector = f"aiq-test={suffix},role=primary"
|
|
|
|
created: list[str] = []
|
|
try:
|
|
ref_a = sandbox_client.create(
|
|
workspace="default",
|
|
name=job_a,
|
|
labels={"aiq-test": suffix, "role": "primary"},
|
|
)
|
|
created.append(ref_a.name)
|
|
ref_b = sandbox_client.create(
|
|
workspace="default",
|
|
name=job_b,
|
|
labels={"aiq-test": suffix, "role": "secondary"},
|
|
)
|
|
created.append(ref_b.name)
|
|
|
|
# Labels round-trip through create and get.
|
|
assert ref_a.labels["role"] == "primary"
|
|
assert (
|
|
dict(sandbox_client.get(job_a, workspace="default").labels)["role"]
|
|
== "primary"
|
|
)
|
|
assert (
|
|
dict(sandbox_client.get(job_b, workspace="default").labels)["role"]
|
|
== "secondary"
|
|
)
|
|
|
|
# A specific selector filters to exactly the primary sandbox.
|
|
assert {
|
|
s.name
|
|
for s in sandbox_client.list_all(
|
|
workspace="default", label_selector=primary_selector
|
|
)
|
|
} == {job_a}
|
|
# The shared group label returns both.
|
|
assert {
|
|
s.name
|
|
for s in sandbox_client.list_all(
|
|
workspace="default", label_selector=group_selector
|
|
)
|
|
} == {
|
|
job_a,
|
|
job_b,
|
|
}
|
|
|
|
# Deleting one removes only it from selector results.
|
|
assert sandbox_client.delete(job_a, workspace="default")
|
|
sandbox_client.wait_deleted(job_a, workspace="default")
|
|
created.remove(job_a)
|
|
assert {
|
|
s.name
|
|
for s in sandbox_client.list_all(
|
|
workspace="default", label_selector=group_selector
|
|
)
|
|
} == {job_b}
|
|
|
|
# Final deletion leaves no matching sandboxes.
|
|
assert sandbox_client.delete(job_b, workspace="default")
|
|
sandbox_client.wait_deleted(job_b, workspace="default")
|
|
created.remove(job_b)
|
|
assert not sandbox_client.list_all(
|
|
workspace="default", label_selector=group_selector
|
|
)
|
|
finally:
|
|
for name in created:
|
|
with contextlib.suppress(Exception):
|
|
sandbox_client.delete(name, workspace="default")
|