tests/test_integration.py
4.3 KB · 119 lines · python Raw
1 """End-to-end integration tests — full handshake + tool call flow."""
2
3 from __future__ import annotations
4
5 import pytest
6 from quantumshield.identity.agent import AgentIdentity
7
8 from pqc_mcp_transport.handshake import PQCHandshake
9 from pqc_mcp_transport.server import PQCMCPServer
10 from pqc_mcp_transport.signer import MessageSigner
11
12
13 @pytest.fixture
14 def integration_server(server_identity: AgentIdentity) -> PQCMCPServer:
15 srv = PQCMCPServer(identity=server_identity, require_auth=True)
16
17 @srv.tool("add", description="Add two numbers")
18 async def add(a: float, b: float) -> float:
19 return a + b
20
21 @srv.tool("echo", description="Echo input")
22 async def echo(message: str) -> str:
23 return message
24
25 return srv
26
27
28 @pytest.mark.asyncio
29 class TestIntegration:
30 async def test_client_server_handshake_and_tool_call(
31 self,
32 integration_server: PQCMCPServer,
33 client_identity: AgentIdentity,
34 server_identity: AgentIdentity,
35 ) -> None:
36 """Full round-trip: handshake -> tool call -> verified response."""
37 # Step 1: Client initiates handshake
38 hs_request, nonce = PQCHandshake.initiate(client_identity)
39
40 # Step 2: Server responds to handshake
41 hs_response_dict = await integration_server.handle_handshake(
42 hs_request.to_dict()
43 )
44
45 # Step 3: Client completes handshake
46 from pqc_mcp_transport.handshake import HandshakeResponse
47
48 hs_response = HandshakeResponse.from_dict(hs_response_dict)
49 session = PQCHandshake.complete(hs_response, client_identity, nonce)
50 assert session.is_valid()
51 assert session.peer_did == server_identity.did
52
53 # Step 4: Client sends a signed tool call
54 client_signer = MessageSigner(client_identity)
55 call_msg = {
56 "jsonrpc": "2.0",
57 "method": "tools/call",
58 "id": "int-1",
59 "params": {"name": "add", "arguments": {"a": 3.0, "b": 4.0}},
60 }
61 signed_call = client_signer.sign_message(call_msg)
62 signed_call["_pqc"]["session_id"] = session.session_id
63
64 # Step 5: Server processes, verifies, and returns signed response
65 response = await integration_server.handle_request(signed_call)
66 assert "_pqc" in response
67
68 # Step 6: Client verifies server response
69 vr = MessageSigner.verify_message(response)
70 assert vr.valid is True
71 assert vr.signer_did == server_identity.did
72
73 stripped = MessageSigner.strip_pqc(response)
74 assert stripped["result"]["content"] == 7.0
75
76 async def test_mutual_authentication(
77 self,
78 integration_server: PQCMCPServer,
79 client_identity: AgentIdentity,
80 server_identity: AgentIdentity,
81 ) -> None:
82 """Both sides verify each other's signatures throughout the flow."""
83 # Handshake
84 hs_request, nonce = PQCHandshake.initiate(client_identity)
85 hs_response_dict = await integration_server.handle_handshake(
86 hs_request.to_dict()
87 )
88 from pqc_mcp_transport.handshake import HandshakeResponse
89
90 hs_response = HandshakeResponse.from_dict(hs_response_dict)
91 session = PQCHandshake.complete(hs_response, client_identity, nonce)
92
93 # Client signs a request — server verifies the client's identity
94 client_signer = MessageSigner(client_identity)
95 call_msg = {
96 "jsonrpc": "2.0",
97 "method": "tools/call",
98 "id": "int-2",
99 "params": {"name": "echo", "arguments": {"message": "mutual-auth-test"}},
100 }
101 signed_call = client_signer.sign_message(call_msg)
102 signed_call["_pqc"]["session_id"] = session.session_id
103
104 # Verify client signature independently
105 client_vr = MessageSigner.verify_message(signed_call)
106 assert client_vr.valid is True
107 assert client_vr.signer_did == client_identity.did
108
109 # Server processes and signs response
110 response = await integration_server.handle_request(signed_call)
111
112 # Verify server signature independently
113 server_vr = MessageSigner.verify_message(response)
114 assert server_vr.valid is True
115 assert server_vr.signer_did == server_identity.did
116
117 stripped = MessageSigner.strip_pqc(response)
118 assert stripped["result"]["content"] == "mutual-auth-test"
119