forked from themanojdesai/python-a2a
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_server.py
More file actions
118 lines (88 loc) · 4.2 KB
/
Copy pathtest_server.py
File metadata and controls
118 lines (88 loc) · 4.2 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
"""
Tests for the server module.
"""
import pytest
from unittest.mock import patch, MagicMock
from python_a2a import (
A2AServer, Message, TextContent, MessageRole, Conversation,
BaseA2AServer, run_server
)
class TestBaseA2AServer:
def test_handle_message_abstract(self):
"""Test that handle_message is abstract"""
with pytest.raises(TypeError):
# Should raise TypeError because handle_message is abstract
server = BaseA2AServer()
def test_handle_conversation(self, echo_server, conversation):
"""Test handling a conversation"""
# Handle the conversation
result = echo_server.handle_conversation(conversation)
# Should add a new message to the conversation
assert len(result.messages) == len(conversation.messages) + 1
# Check the new message
new_message = result.messages[-1]
assert new_message.role == MessageRole.AGENT
assert new_message.parent_message_id == conversation.messages[-1].message_id
assert new_message.conversation_id == conversation.conversation_id
class TestA2AServer:
def test_init_with_handler(self):
"""Test initializing with a custom handler"""
handler = lambda message: Message(
content=TextContent(text="Custom response"),
role=MessageRole.AGENT
)
server = A2AServer(message_handler=handler)
assert server.message_handler == handler
def test_handle_message_with_handler(self, text_message):
"""Test handling a message with a custom handler"""
# Create a mock handler
handler = MagicMock(return_value=Message(
content=TextContent(text="Custom response"),
role=MessageRole.AGENT
))
# Create server with the handler
server = A2AServer(message_handler=handler)
# Handle the message
response = server.handle_message(text_message)
# Check that the handler was called
handler.assert_called_once_with(text_message)
# Check the response
assert response.content.type == "text"
assert response.content.text == "Custom response"
assert response.role == MessageRole.AGENT
def test_handle_message_default(self, text_message):
"""Test default message handling"""
server = A2AServer()
response = server.handle_message(text_message)
assert response.content.type == "text"
assert response.content.text.startswith("Echo:")
assert response.role == MessageRole.AGENT
def test_get_metadata(self):
"""Test getting server metadata"""
server = A2AServer()
metadata = server.get_metadata()
assert "agent_type" in metadata
assert "capabilities" in metadata
assert "text" in metadata["capabilities"]
class TestRunServer:
def test_create_flask_app(self, echo_server):
"""Test creating a Flask app"""
with patch('python_a2a.server.http.Flask') as MockFlask:
MockFlask.return_value = MagicMock()
from python_a2a.server.http import create_flask_app
app = create_flask_app(echo_server)
# Check that Flask was initialized
MockFlask.assert_called_once()
# The app should have routes for /a2a, /a2a/metadata, and /a2a/health
app.route.assert_any_call("/a2a", methods=["POST"])
app.route.assert_any_call("/a2a/metadata", methods=["GET"])
app.route.assert_any_call("/a2a/health", methods=["GET"])
def test_run_server(self, echo_server):
"""Test running a server"""
with patch('python_a2a.server.http.create_flask_app') as mock_create_app:
mock_app = MagicMock()
mock_create_app.return_value = mock_app
run_server(echo_server, host="localhost", port=8080, debug=True)
# Check that the app was created and run
mock_create_app.assert_called_once_with(echo_server)
mock_app.run.assert_called_once_with(host="localhost", port=8080, debug=True)