flymcprust_core/
testing.rs

1//! In-process test client (feature `test-client`).
2//!
3//! [`TestClient`] talks to a [`Server`] without any I/O, but through the
4//! full serialization path: every request is encoded to bytes, parsed by the
5//! server exactly as a transport would, and the response is decoded from
6//! bytes. Tests therefore exercise the real wire format.
7//!
8//! ```
9//! # use flymcprust_core::{Server, testing::TestClient};
10//! # async fn demo(server: Server) -> Result<(), flymcprust_core::testing::ClientError> {
11//! let client = TestClient::new(server);
12//! let tools = client.list_tools().await?;
13//! let result = client.call_tool("add", serde_json::json!({ "a": 1, "b": 2 })).await?;
14//! # Ok(()) }
15//! ```
16
17use std::sync::atomic::{AtomicI64, Ordering};
18
19use serde::de::DeserializeOwned;
20use serde_json::{Value, json};
21use tokio_util::sync::CancellationToken;
22
23use crate::context::TransportContext;
24use crate::server::Server;
25use crate::spec::{self, methods};
26
27/// An error returned by [`TestClient`].
28#[derive(Debug, thiserror::Error)]
29#[non_exhaustive]
30pub enum ClientError {
31    /// The server answered with a JSON-RPC error.
32    #[error("JSON-RPC error {code}: {message}")]
33    Rpc {
34        /// The error code.
35        code: i64,
36        /// The error message.
37        message: String,
38        /// The optional `data` member.
39        data: Option<Value>,
40    },
41    /// The server sent no response (e.g. the request was a notification).
42    #[error("the server sent no response")]
43    NoResponse,
44    /// The response could not be decoded.
45    #[error("invalid response: {0}")]
46    Decode(#[from] serde_json::Error),
47}
48
49/// A client that calls a [`Server`] in process.
50#[derive(Debug)]
51pub struct TestClient {
52    server: Server,
53    next_id: AtomicI64,
54    meta: Value,
55}
56
57impl TestClient {
58    /// Creates a client that declares no capabilities and the server's
59    /// protocol version.
60    #[must_use]
61    pub fn new(server: Server) -> Self {
62        let meta = json!({
63            "io.modelcontextprotocol/protocolVersion": spec::PROTOCOL_VERSION,
64            "io.modelcontextprotocol/clientCapabilities": {},
65            "io.modelcontextprotocol/clientInfo": { "name": "flymcprust-test-client", "version": env!("CARGO_PKG_VERSION") },
66        });
67        Self {
68            server,
69            next_id: AtomicI64::new(1),
70            meta,
71        }
72    }
73
74    /// Replaces the `_meta` sent with every request.
75    #[must_use]
76    pub fn with_meta(mut self, meta: Value) -> Self {
77        self.meta = meta;
78        self
79    }
80
81    /// Calls `server/discover`.
82    ///
83    /// # Errors
84    ///
85    /// Returns a [`ClientError`] for JSON-RPC errors or undecodable results.
86    pub async fn discover(&self) -> Result<spec::DiscoverResult, ClientError> {
87        self.request(methods::SERVER_DISCOVER, json!({})).await
88    }
89
90    /// Calls `tools/list`.
91    ///
92    /// # Errors
93    ///
94    /// Returns a [`ClientError`] for JSON-RPC errors or undecodable results.
95    pub async fn list_tools(&self) -> Result<spec::ListToolsResult, ClientError> {
96        self.request(methods::TOOLS_LIST, json!({})).await
97    }
98
99    /// Calls `tools/call`. Tool failures are results with `is_error`, not
100    /// errors.
101    ///
102    /// # Errors
103    ///
104    /// Returns a [`ClientError`] for JSON-RPC errors or undecodable results.
105    pub async fn call_tool(
106        &self,
107        name: &str,
108        arguments: Value,
109    ) -> Result<spec::CallToolResult, ClientError> {
110        self.request(
111            methods::TOOLS_CALL,
112            json!({ "name": name, "arguments": arguments }),
113        )
114        .await
115    }
116
117    /// Calls `resources/list`.
118    ///
119    /// # Errors
120    ///
121    /// Returns a [`ClientError`] for JSON-RPC errors or undecodable results.
122    pub async fn list_resources(&self) -> Result<spec::ListResourcesResult, ClientError> {
123        self.request(methods::RESOURCES_LIST, json!({})).await
124    }
125
126    /// Calls `resources/templates/list`.
127    ///
128    /// # Errors
129    ///
130    /// Returns a [`ClientError`] for JSON-RPC errors or undecodable results.
131    pub async fn list_resource_templates(
132        &self,
133    ) -> Result<spec::ListResourceTemplatesResult, ClientError> {
134        self.request(methods::RESOURCES_TEMPLATES_LIST, json!({}))
135            .await
136    }
137
138    /// Calls `resources/read`.
139    ///
140    /// # Errors
141    ///
142    /// Returns a [`ClientError`] for JSON-RPC errors (including
143    /// resource-not-found) or undecodable results.
144    pub async fn read_resource(&self, uri: &str) -> Result<spec::ReadResourceResult, ClientError> {
145        self.request(methods::RESOURCES_READ, json!({ "uri": uri }))
146            .await
147    }
148
149    /// Calls `prompts/list`.
150    ///
151    /// # Errors
152    ///
153    /// Returns a [`ClientError`] for JSON-RPC errors or undecodable results.
154    pub async fn list_prompts(&self) -> Result<spec::ListPromptsResult, ClientError> {
155        self.request(methods::PROMPTS_LIST, json!({})).await
156    }
157
158    /// Calls `prompts/get` with string arguments.
159    ///
160    /// # Errors
161    ///
162    /// Returns a [`ClientError`] for JSON-RPC errors or undecodable results.
163    pub async fn get_prompt(
164        &self,
165        name: &str,
166        arguments: Value,
167    ) -> Result<spec::GetPromptResult, ClientError> {
168        self.request(
169            methods::PROMPTS_GET,
170            json!({ "name": name, "arguments": arguments }),
171        )
172        .await
173    }
174
175    /// Calls `completion/complete` for a prompt argument.
176    ///
177    /// # Errors
178    ///
179    /// Returns a [`ClientError`] for JSON-RPC errors or undecodable results.
180    pub async fn complete_prompt(
181        &self,
182        prompt: &str,
183        argument: &str,
184        value: &str,
185    ) -> Result<spec::CompleteResult, ClientError> {
186        self.request(
187            methods::COMPLETION_COMPLETE,
188            json!({
189                "ref": { "type": "ref/prompt", "name": prompt },
190                "argument": { "name": argument, "value": value },
191            }),
192        )
193        .await
194    }
195
196    /// Calls `completion/complete` for a resource-template variable.
197    ///
198    /// # Errors
199    ///
200    /// Returns a [`ClientError`] for JSON-RPC errors or undecodable results.
201    pub async fn complete_resource(
202        &self,
203        uri_template: &str,
204        variable: &str,
205        value: &str,
206    ) -> Result<spec::CompleteResult, ClientError> {
207        self.request(
208            methods::COMPLETION_COMPLETE,
209            json!({
210                "ref": { "type": "ref/resource", "uri": uri_template },
211                "argument": { "name": variable, "value": value },
212            }),
213        )
214        .await
215    }
216
217    /// Sends any request; `_meta` is added to `params` unless present.
218    ///
219    /// # Errors
220    ///
221    /// Returns a [`ClientError`] for JSON-RPC errors or undecodable results.
222    pub async fn request<T: DeserializeOwned>(
223        &self,
224        method: &str,
225        mut params: Value,
226    ) -> Result<T, ClientError> {
227        if let Value::Object(map) = &mut params {
228            map.entry("_meta").or_insert_with(|| self.meta.clone());
229        }
230        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
231        let message = json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params });
232        let bytes = serde_json::to_vec(&message)?;
233        let response = self
234            .server
235            .handle_bytes(
236                &bytes,
237                TransportContext::default(),
238                CancellationToken::new(),
239            )
240            .await
241            .ok_or(ClientError::NoResponse)?;
242        let mut envelope: Value = serde_json::from_slice(&response.to_vec())?;
243        if let Some(error) = envelope.get_mut("error").map(Value::take) {
244            let error: spec::Error = serde_json::from_value(error)?;
245            return Err(ClientError::Rpc {
246                code: error.code,
247                message: error.message,
248                data: error.data,
249            });
250        }
251        let result = envelope
252            .get_mut("result")
253            .map(Value::take)
254            .unwrap_or_default();
255        Ok(serde_json::from_value(result)?)
256    }
257}
flymcprust_core/testing.rs

flymcprust_core/
testing.rs
