wehub-resource-sync
442 lines
14 KiB
Python
442 lines
14 KiB
Python
import json
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from mlflow.assistant.providers.base import clear_config_cache
|
|
from mlflow.assistant.providers.codex import CodexProvider
|
|
from mlflow.assistant.types import EventType
|
|
|
|
|
|
def _make_stdout_lines(*dicts) -> list[bytes]:
|
|
return [json.dumps(d).encode() + b"\n" for d in dicts]
|
|
|
|
|
|
def _mock_process(stdout_lines=None, returncode=0, stderr=b""):
|
|
"""Create a mock process with async stdout iteration and stdin support."""
|
|
process = MagicMock()
|
|
process.returncode = returncode
|
|
|
|
# Mock stdin (write, drain, close, wait_closed)
|
|
process.stdin = MagicMock()
|
|
process.stdin.write = MagicMock()
|
|
process.stdin.drain = AsyncMock()
|
|
process.stdin.close = MagicMock()
|
|
process.stdin.wait_closed = AsyncMock()
|
|
|
|
# Mock stdout as an async iterator
|
|
async def _aiter():
|
|
for line in stdout_lines or []:
|
|
yield line
|
|
|
|
process.stdout = _aiter()
|
|
|
|
# Mock stderr.read()
|
|
process.stderr = MagicMock()
|
|
process.stderr.read = AsyncMock(return_value=stderr)
|
|
|
|
# Mock wait
|
|
process.wait = AsyncMock()
|
|
process.kill = MagicMock()
|
|
|
|
return process
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def config(tmp_path):
|
|
config_file = tmp_path / "config.json"
|
|
config_file.write_text('{"providers": {"codex": {"model": "default"}}}')
|
|
clear_config_cache()
|
|
with patch("mlflow.assistant.config.CONFIG_PATH", config_file):
|
|
yield config_file
|
|
clear_config_cache()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("which_return", "expected"),
|
|
[
|
|
("/usr/local/bin/codex", True),
|
|
(None, False),
|
|
],
|
|
)
|
|
def test_is_available(which_return, expected):
|
|
with patch("mlflow.assistant.providers.codex.shutil.which", return_value=which_return):
|
|
provider = CodexProvider()
|
|
assert provider.is_available() is expected
|
|
|
|
|
|
def test_provider_name():
|
|
assert CodexProvider().name == "codex"
|
|
|
|
|
|
def test_provider_display_name():
|
|
assert CodexProvider().display_name == "OpenAI Codex"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_yields_error_when_codex_not_found():
|
|
with patch("mlflow.assistant.providers.codex.shutil.which", return_value=None):
|
|
provider = CodexProvider()
|
|
events = [e async for e in provider.astream("test prompt", "http://localhost:5000")]
|
|
|
|
assert len(events) == 1
|
|
assert events[0].type == EventType.ERROR
|
|
assert "codex" in events[0].data["error"].lower()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_yields_agent_message_text():
|
|
stdout_lines = _make_stdout_lines(
|
|
{"type": "thread.started", "thread_id": "t1"},
|
|
{"type": "turn.started"},
|
|
{
|
|
"type": "item.completed",
|
|
"item": {"id": "i1", "type": "agent_message", "text": "Hello there!"},
|
|
},
|
|
{"type": "turn.completed", "usage": {"input_tokens": 10, "output_tokens": 5}},
|
|
)
|
|
|
|
mock_proc = _mock_process(stdout_lines=stdout_lines)
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
),
|
|
):
|
|
provider = CodexProvider()
|
|
events = [e async for e in provider.astream("hi", "http://localhost:5000")]
|
|
|
|
message_events = [e for e in events if e.type == EventType.MESSAGE]
|
|
assert len(message_events) == 1
|
|
assert message_events[0].data["message"]["content"][0]["text"] == "Hello there!"
|
|
|
|
done_events = [e for e in events if e.type == EventType.DONE]
|
|
assert len(done_events) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_emits_usage_event_from_turn_completed():
|
|
stdout_lines = _make_stdout_lines(
|
|
{"type": "thread.started", "thread_id": "t1"},
|
|
{"type": "item.completed", "item": {"type": "agent_message", "text": "done"}},
|
|
{
|
|
"type": "turn.completed",
|
|
"usage": {"input_tokens": 10, "cached_input_tokens": 4, "output_tokens": 5},
|
|
},
|
|
)
|
|
|
|
mock_proc = _mock_process(stdout_lines=stdout_lines)
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.calculate_cost_by_model_and_token_usage",
|
|
return_value=None,
|
|
) as mock_cost,
|
|
):
|
|
provider = CodexProvider()
|
|
events = [e async for e in provider.astream("hi", "http://localhost:5000")]
|
|
|
|
usage_events = [
|
|
e
|
|
for e in events
|
|
if e.type == EventType.STREAM_EVENT and e.data["event"].get("type") == "usage"
|
|
]
|
|
assert len(usage_events) == 1
|
|
usage = usage_events[0].data["event"]["usage"]
|
|
assert usage["prompt_tokens"] == 10
|
|
assert usage["completion_tokens"] == 5
|
|
assert usage["total_tokens"] == 15
|
|
assert usage["total_cost_usd"] is None
|
|
mock_cost.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_ignores_non_agent_message_items():
|
|
mcp_item = {"id": "i1", "type": "mcp_tool_call", "text": "ignored"}
|
|
agent_item = {"id": "i2", "type": "agent_message", "text": "kept"}
|
|
stdout_lines = _make_stdout_lines(
|
|
{"type": "item.completed", "item": mcp_item},
|
|
{"type": "item.completed", "item": agent_item},
|
|
)
|
|
|
|
mock_proc = _mock_process(stdout_lines=stdout_lines)
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
),
|
|
):
|
|
provider = CodexProvider()
|
|
events = [e async for e in provider.astream("hi", "http://localhost:5000")]
|
|
|
|
message_events = [e for e in events if e.type == EventType.MESSAGE]
|
|
assert len(message_events) == 1
|
|
assert message_events[0].data["message"]["content"][0]["text"] == "kept"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_yields_error_on_nonzero_exit():
|
|
mock_proc = _mock_process(returncode=1, stderr=b"OPENAI_API_KEY not set")
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
),
|
|
):
|
|
provider = CodexProvider()
|
|
events = [e async for e in provider.astream("hi", "http://localhost:5000")]
|
|
|
|
error_events = [e for e in events if e.type == EventType.ERROR]
|
|
assert len(error_events) == 1
|
|
assert "OPENAI_API_KEY" in error_events[0].data["error"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_yields_interrupted_on_sigkill():
|
|
mock_proc = _mock_process(returncode=-9)
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
),
|
|
):
|
|
provider = CodexProvider()
|
|
events = [e async for e in provider.astream("hi", "http://localhost:5000")]
|
|
|
|
assert len(events) == 1
|
|
assert events[0].type == EventType.INTERRUPTED
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_builds_correct_command():
|
|
mock_proc = _mock_process()
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
) as mock_exec,
|
|
):
|
|
provider = CodexProvider()
|
|
_ = [e async for e in provider.astream("test prompt", "http://localhost:5000")]
|
|
|
|
args = mock_exec.call_args[0]
|
|
assert "/usr/bin/codex" in args
|
|
assert "exec" in args
|
|
assert "--json" in args
|
|
assert "--sandbox" in args
|
|
assert "danger-full-access" in args
|
|
assert "--dangerously-bypass-approvals-and-sandbox" not in args
|
|
assert "--ephemeral" not in args
|
|
assert "--skip-git-repo-check" in args
|
|
assert args[-1] == "-"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_includes_model_flag_when_configured(tmp_path):
|
|
config_file = tmp_path / "config.json"
|
|
config_file.write_text('{"providers": {"codex": {"model": "o4-mini"}}}')
|
|
|
|
mock_proc = _mock_process()
|
|
|
|
clear_config_cache()
|
|
with (
|
|
patch("mlflow.assistant.config.CONFIG_PATH", config_file),
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
) as mock_exec,
|
|
):
|
|
provider = CodexProvider()
|
|
_ = [e async for e in provider.astream("prompt", "http://localhost:5000")]
|
|
|
|
args = mock_exec.call_args[0]
|
|
assert "-m" in args
|
|
assert "o4-mini" in args
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_skips_model_flag_when_default(tmp_path):
|
|
config_file = tmp_path / "config.json"
|
|
config_file.write_text('{"providers": {"codex": {"model": "default"}}}')
|
|
|
|
mock_proc = _mock_process()
|
|
|
|
clear_config_cache()
|
|
with (
|
|
patch("mlflow.assistant.config.CONFIG_PATH", config_file),
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
) as mock_exec,
|
|
):
|
|
provider = CodexProvider()
|
|
_ = [e async for e in provider.astream("prompt", "http://localhost:5000")]
|
|
|
|
args = mock_exec.call_args[0]
|
|
assert "-m" not in args
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_sends_prompt_via_stdin():
|
|
mock_proc = _mock_process()
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
),
|
|
):
|
|
provider = CodexProvider()
|
|
_ = [e async for e in provider.astream("my question", "http://localhost:5000")]
|
|
|
|
stdin_bytes = mock_proc.stdin.write.call_args[0][0]
|
|
assert b"<system_instructions>" in stdin_bytes
|
|
assert b"my question" in stdin_bytes
|
|
mock_proc.stdin.drain.assert_awaited_once()
|
|
mock_proc.stdin.close.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_omits_system_instructions_on_resume():
|
|
mock_proc = _mock_process()
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
),
|
|
):
|
|
provider = CodexProvider()
|
|
_ = [
|
|
e
|
|
async for e in provider.astream(
|
|
"follow up", "http://localhost:5000", session_id="abc-123"
|
|
)
|
|
]
|
|
|
|
stdin_bytes = mock_proc.stdin.write.call_args[0][0]
|
|
assert b"<system_instructions>" not in stdin_bytes
|
|
assert b"follow up" in stdin_bytes
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_captures_session_id_from_thread_started():
|
|
stdout_lines = _make_stdout_lines(
|
|
{"type": "thread.started", "thread_id": "abc-123"},
|
|
{"type": "item.completed", "item": {"type": "agent_message", "text": "hi"}},
|
|
)
|
|
mock_proc = _mock_process(stdout_lines=stdout_lines)
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
),
|
|
):
|
|
provider = CodexProvider()
|
|
events = [e async for e in provider.astream("hi", "http://localhost:5000")]
|
|
|
|
done_events = [e for e in events if e.type == EventType.DONE]
|
|
assert len(done_events) == 1
|
|
assert done_events[0].data["session_id"] == "abc-123"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_resumes_session_when_session_id_provided():
|
|
mock_proc = _mock_process()
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
) as mock_exec,
|
|
):
|
|
provider = CodexProvider()
|
|
_ = [
|
|
e
|
|
async for e in provider.astream(
|
|
"follow up", "http://localhost:5000", session_id="abc-123"
|
|
)
|
|
]
|
|
|
|
args = mock_exec.call_args[0]
|
|
assert "resume" in args
|
|
assert "abc-123" in args
|
|
resume_idx = list(args).index("resume")
|
|
assert args[resume_idx + 1] == "abc-123"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_saves_and_clears_process_pid():
|
|
mock_proc = _mock_process()
|
|
mock_proc.pid = 12345
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
),
|
|
patch("mlflow.assistant.providers.codex.save_process_pid") as mock_save,
|
|
patch("mlflow.assistant.providers.codex.clear_process_pid") as mock_clear,
|
|
):
|
|
provider = CodexProvider()
|
|
_ = [
|
|
e
|
|
async for e in provider.astream(
|
|
"hi", "http://localhost:5000", mlflow_session_id="session-xyz"
|
|
)
|
|
]
|
|
|
|
mock_save.assert_called_once_with("session-xyz", 12345)
|
|
mock_clear.assert_called_once_with("session-xyz")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_astream_ignores_invalid_json_lines():
|
|
valid_item = {
|
|
"type": "item.completed",
|
|
"item": {"type": "agent_message", "text": "valid"},
|
|
}
|
|
stdout_lines = [
|
|
b"not json\n",
|
|
json.dumps(valid_item).encode() + b"\n",
|
|
]
|
|
|
|
mock_proc = _mock_process(stdout_lines=stdout_lines)
|
|
|
|
with (
|
|
patch("mlflow.assistant.providers.codex.shutil.which", return_value="/usr/bin/codex"),
|
|
patch(
|
|
"mlflow.assistant.providers.codex.asyncio.create_subprocess_exec",
|
|
return_value=mock_proc,
|
|
),
|
|
):
|
|
provider = CodexProvider()
|
|
events = [e async for e in provider.astream("hi", "http://localhost:5000")]
|
|
|
|
message_events = [e for e in events if e.type == EventType.MESSAGE]
|
|
assert len(message_events) == 1
|
|
assert message_events[0].data["message"]["content"][0]["text"] == "valid"
|