79 lines
No EOL
2.5 KiB
Python
79 lines
No EOL
2.5 KiB
Python
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from lambda_agent_api.server import *
|
|
from lambda_agent_api.client import *
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"data, expected_type",
|
|
[
|
|
({"type": "USER_MESSAGE", "text": "hello"}, MsgUserMessage),
|
|
],
|
|
)
|
|
def test_client_message_valid(data, expected_type):
|
|
msg = ClientMessage.validate_python(data)
|
|
assert isinstance(msg, expected_type)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"data",
|
|
[
|
|
{"type": "UNKNOWN", "text": "hello"},
|
|
{"type": "USER_MESSAGE"}, # нет text
|
|
],
|
|
)
|
|
def test_client_message_invalid(data):
|
|
with pytest.raises(ValidationError):
|
|
ClientMessage.validate_python(data)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"data, expected_type",
|
|
[
|
|
({"type": "STATUS"}, MsgStatus),
|
|
({"type": "AGENT_EVENT_TEXT_CHUNK", "text": "hi"}, MsgEventTextChunk),
|
|
({"type": "AGENT_EVENT_TOOL_CALL_CHUNK", "tool_name": "search", "args_chunk": "{\"q\": \"hello\"}"}, MsgEventToolCallChunk),
|
|
({"type": "AGENT_EVENT_TOOL_RESULT", "tool_name": "search", "result": {"items": [1, 2, 3]}}, MsgEventToolResult),
|
|
({"type": "AGENT_EVENT_CUSTOM_UPDATE", "payload": {"status": "in_progress", "progress": 50}}, MsgEventCustomUpdate),
|
|
({"type": "AGENT_EVENT_END", "tokens_used": 10}, MsgEventEnd),
|
|
({"type": "ERROR", "code": "E1", "details": "fail"}, MsgError),
|
|
({"type": "GRACEFUL_DISCONNECT"}, MsgGracefulDisconnect),
|
|
],
|
|
)
|
|
def test_server_message_valid(data, expected_type):
|
|
msg = ServerMessage.validate_python(data)
|
|
assert isinstance(msg, expected_type)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"data",
|
|
[
|
|
{"type": "AGENT_EVENT_TEXT_CHUNK"}, # нет text
|
|
{"type": "AGENT_EVENT_TOOL_RESULT", "tool_name": "search"}, # нет result
|
|
{"type": "AGENT_EVENT_CUSTOM_UPDATE"}, # нет payload
|
|
{"type": "AGENT_EVENT_END"}, # нет tokens_used
|
|
{"type": "ERROR", "code": "E1"}, # нет details
|
|
{"type": "UNKNOWN"},
|
|
],
|
|
)
|
|
def test_server_message_invalid(data):
|
|
with pytest.raises(ValidationError):
|
|
ServerMessage.validate_python(data)
|
|
|
|
|
|
def test_validate_json():
|
|
json_data = '{"type": "AGENT_EVENT_TEXT_CHUNK", "text": "hello"}'
|
|
msg = ServerMessage.validate_json(json_data)
|
|
|
|
assert isinstance(msg, MsgEventTextChunk)
|
|
assert msg.text == "hello"
|
|
|
|
|
|
def test_model_dump_roundtrip():
|
|
data = {"type": "ERROR", "code": "E1", "details": "fail"}
|
|
|
|
msg = ServerMessage.validate_python(data)
|
|
dumped = msg.model_dump()
|
|
|
|
assert dumped == data |