wehub-resource-sync
123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
import importlib.util
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from mlflow.assistant.providers.claude_code import ClaudeCodeProvider
|
|
from mlflow.assistant.types import EventType
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
DEV_STUBS = REPO_ROOT / "dev" / "dev_stubs"
|
|
CLAUDE_CLI = DEV_STUBS / "claude_cli.py"
|
|
|
|
|
|
def _load(name: str, path: Path):
|
|
# dev/ is not an installed package; load this module directly by path. It only
|
|
# imports stdlib and references siblings as files, so standalone load works.
|
|
spec = importlib.util.spec_from_file_location(name, path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
# Register before exec so @dataclass can resolve the module's globals by name.
|
|
sys.modules[name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
dev_stubs = _load("mlflow_dev_stubs", DEV_STUBS / "__init__.py")
|
|
|
|
|
|
# --- claude CLI stub ---------------------------------------------------------
|
|
|
|
|
|
def run_claude(*args: str) -> subprocess.CompletedProcess[str]:
|
|
return subprocess.run(
|
|
[sys.executable, CLAUDE_CLI, *args], capture_output=True, text=True, timeout=30
|
|
)
|
|
|
|
|
|
def test_claude_auth_probe_exits_zero_with_valid_json():
|
|
result = run_claude("-p", "hi", "--max-turns", "1", "--output-format", "json")
|
|
assert result.returncode == 0
|
|
payload = json.loads(result.stdout)
|
|
assert payload["type"] == "result"
|
|
assert payload["is_error"] is False
|
|
|
|
|
|
def test_claude_stream_json_parses_into_message_and_done_events():
|
|
result = run_claude(
|
|
"-p",
|
|
"hello",
|
|
"--output-format",
|
|
"stream-json",
|
|
"--verbose",
|
|
"--append-system-prompt",
|
|
"ignored prompt with --output-format text and --resume decoys",
|
|
)
|
|
assert result.returncode == 0
|
|
|
|
provider = ClaudeCodeProvider()
|
|
events = [
|
|
event
|
|
for line in result.stdout.splitlines()
|
|
if line.strip()
|
|
if (event := provider._parse_message_to_event(json.loads(line))) is not None
|
|
]
|
|
assert [e.type for e in events] == [EventType.MESSAGE, EventType.DONE]
|
|
assert events[0].data["message"]["content"][0]["text"]
|
|
assert events[1].data["session_id"]
|
|
|
|
|
|
def test_claude_resume_reuses_session_id():
|
|
result = run_claude("-p", "hi", "--output-format", "stream-json", "--resume", "sess-abc")
|
|
assert result.returncode == 0
|
|
events = [json.loads(line) for line in result.stdout.splitlines() if line.strip()]
|
|
assert {e["session_id"] for e in events if "session_id" in e} == {"sess-abc"}
|
|
|
|
|
|
# --- stub registry -----------------------------------------------------------
|
|
|
|
|
|
def test_install_stubs_rejects_unknown_name():
|
|
with pytest.raises(ValueError, match="Unknown stub"):
|
|
dev_stubs.install_stubs(["not-a-stub"])
|
|
|
|
|
|
def test_install_stubs_cleans_up_temp_dirs_on_partial_failure(monkeypatch):
|
|
tmp_root = Path(tempfile.gettempdir())
|
|
before = set(tmp_root.glob("mlflow-dev-stub-bin-*"))
|
|
|
|
def boom(_result):
|
|
raise RuntimeError("boom")
|
|
|
|
monkeypatch.setitem(dev_stubs._INSTALLERS, "boom", boom)
|
|
with pytest.raises(RuntimeError, match="boom"):
|
|
# claude stages a temp dir, then `boom` fails -- the dir must not leak.
|
|
dev_stubs.install_stubs(["claude", "boom"])
|
|
|
|
assert set(tmp_root.glob("mlflow-dev-stub-bin-*")) == before
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform == "win32",
|
|
reason="the claude shim is a POSIX shell script; run_dev_server (its consumer) is POSIX-only",
|
|
)
|
|
def test_install_claude_stages_working_shim():
|
|
result = dev_stubs.install_stubs(["claude"])
|
|
try:
|
|
(shim_dir,) = result.path_prepend
|
|
probe = subprocess.run(
|
|
[shim_dir / "claude", "-p", "hi", "--output-format", "json"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=30,
|
|
)
|
|
assert probe.returncode == 0
|
|
assert json.loads(probe.stdout)["type"] == "result"
|
|
finally:
|
|
for path in result.cleanup_paths:
|
|
shutil.rmtree(path, ignore_errors=True)
|