mcpkit_testing/
mock.rs

1//! Mock implementations for testing.
2//!
3//! This module provides mock servers and tools that can be used in unit tests.
4//! The mocks are designed to be flexible and configurable.
5
6use mcpkit_core::capability::{ServerCapabilities, ServerInfo};
7use mcpkit_core::error::McpError;
8use mcpkit_core::types::{
9    Content, GetPromptResult, Prompt, PromptMessage, Resource, ResourceContents, Tool,
10    ToolAnnotations, ToolOutput,
11};
12use mcpkit_server::{Context, PromptHandler, ResourceHandler, ServerHandler, ToolHandler};
13use serde_json::Value;
14use std::collections::HashMap;
15use std::future::Future;
16use std::sync::Arc;
17
18/// A mock tool with configurable behavior.
19pub struct MockTool {
20    /// Tool name.
21    pub name: String,
22    /// Tool description.
23    pub description: Option<String>,
24    /// Input schema.
25    pub input_schema: Value,
26    /// Annotations.
27    pub annotations: Option<ToolAnnotations>,
28    /// Response to return.
29    pub response: MockResponse,
30}
31
32/// Type of response a mock tool should return.
33#[derive(Clone)]
34pub enum MockResponse {
35    /// Return a successful text response.
36    Text(String),
37    /// Return a successful JSON response.
38    Json(Value),
39    /// Return an error.
40    Error(String),
41    /// Return a dynamic response based on input.
42    Dynamic(Arc<dyn Fn(Value) -> Result<ToolOutput, McpError> + Send + Sync>),
43}
44
45impl MockTool {
46    /// Create a new mock tool.
47    pub fn new(name: impl Into<String>) -> Self {
48        Self {
49            name: name.into(),
50            description: None,
51            input_schema: serde_json::json!({
52                "type": "object",
53                "properties": {}
54            }),
55            annotations: None,
56            response: MockResponse::Text("OK".to_string()),
57        }
58    }
59
60    /// Set the description.
61    pub fn description(mut self, description: impl Into<String>) -> Self {
62        self.description = Some(description.into());
63        self
64    }
65
66    /// Set the input schema.
67    #[must_use]
68    pub fn input_schema(mut self, schema: Value) -> Self {
69        self.input_schema = schema;
70        self
71    }
72
73    /// Set annotations.
74    #[must_use]
75    pub fn annotations(mut self, annotations: ToolAnnotations) -> Self {
76        self.annotations = Some(annotations);
77        self
78    }
79
80    /// Set the tool to return a text response.
81    pub fn returns_text(mut self, text: impl Into<String>) -> Self {
82        self.response = MockResponse::Text(text.into());
83        self
84    }
85
86    /// Set the tool to return a JSON response.
87    #[must_use]
88    pub fn returns_json(mut self, json: Value) -> Self {
89        self.response = MockResponse::Json(json);
90        self
91    }
92
93    /// Set the tool to return an error.
94    pub fn returns_error(mut self, message: impl Into<String>) -> Self {
95        self.response = MockResponse::Error(message.into());
96        self
97    }
98
99    /// Set a dynamic handler.
100    pub fn handler<F>(mut self, handler: F) -> Self
101    where
102        F: Fn(Value) -> Result<ToolOutput, McpError> + Send + Sync + 'static,
103    {
104        self.response = MockResponse::Dynamic(Arc::new(handler));
105        self
106    }
107
108    /// Convert to a Tool definition.
109    #[must_use]
110    pub fn to_tool(&self) -> Tool {
111        Tool {
112            name: self.name.clone(),
113            description: self.description.clone(),
114            input_schema: self.input_schema.clone(),
115            annotations: self.annotations.clone(),
116        }
117    }
118
119    /// Call the tool.
120    pub fn call(&self, args: Value) -> Result<ToolOutput, McpError> {
121        match &self.response {
122            MockResponse::Text(text) => Ok(ToolOutput::text(text.clone())),
123            MockResponse::Json(json) => Ok(ToolOutput::text(serde_json::to_string_pretty(json)?)),
124            MockResponse::Error(msg) => Ok(ToolOutput::error(msg.clone())),
125            MockResponse::Dynamic(f) => f(args),
126        }
127    }
128}
129
130/// A mock resource.
131pub struct MockResource {
132    /// Resource URI.
133    pub uri: String,
134    /// Resource name.
135    pub name: String,
136    /// Resource description.
137    pub description: Option<String>,
138    /// MIME type.
139    pub mime_type: Option<String>,
140    /// Resource content.
141    pub content: String,
142}
143
144impl MockResource {
145    /// Create a new mock resource.
146    pub fn new(uri: impl Into<String>, name: impl Into<String>) -> Self {
147        Self {
148            uri: uri.into(),
149            name: name.into(),
150            description: None,
151            mime_type: None,
152            content: String::new(),
153        }
154    }
155
156    /// Set the description.
157    pub fn description(mut self, description: impl Into<String>) -> Self {
158        self.description = Some(description.into());
159        self
160    }
161
162    /// Set the MIME type.
163    pub fn mime_type(mut self, mime_type: impl Into<String>) -> Self {
164        self.mime_type = Some(mime_type.into());
165        self
166    }
167
168    /// Set the content.
169    pub fn content(mut self, content: impl Into<String>) -> Self {
170        self.content = content.into();
171        self
172    }
173
174    /// Convert to a Resource definition.
175    #[must_use]
176    pub fn to_resource(&self) -> Resource {
177        Resource {
178            uri: self.uri.clone(),
179            name: self.name.clone(),
180            description: self.description.clone(),
181            mime_type: self.mime_type.clone(),
182            size: Some(self.content.len() as u64),
183            annotations: None,
184        }
185    }
186
187    /// Get the resource contents.
188    #[must_use]
189    pub fn to_contents(&self) -> ResourceContents {
190        ResourceContents {
191            uri: self.uri.clone(),
192            mime_type: self.mime_type.clone(),
193            text: Some(self.content.clone()),
194            blob: None,
195        }
196    }
197}
198
199/// A mock prompt.
200pub struct MockPrompt {
201    /// Prompt name.
202    pub name: String,
203    /// Prompt description.
204    pub description: Option<String>,
205    /// Message template.
206    pub template: String,
207}
208
209impl MockPrompt {
210    /// Create a new mock prompt.
211    pub fn new(name: impl Into<String>) -> Self {
212        Self {
213            name: name.into(),
214            description: None,
215            template: String::new(),
216        }
217    }
218
219    /// Set the description.
220    pub fn description(mut self, description: impl Into<String>) -> Self {
221        self.description = Some(description.into());
222        self
223    }
224
225    /// Set the message template.
226    pub fn template(mut self, template: impl Into<String>) -> Self {
227        self.template = template.into();
228        self
229    }
230
231    /// Convert to a Prompt definition.
232    #[must_use]
233    pub fn to_prompt(&self) -> Prompt {
234        Prompt {
235            name: self.name.clone(),
236            description: self.description.clone(),
237            arguments: None,
238        }
239    }
240}
241
242/// Builder for constructing mock servers.
243pub struct MockServerBuilder {
244    name: String,
245    version: String,
246    tools: Vec<MockTool>,
247    resources: Vec<MockResource>,
248    prompts: Vec<MockPrompt>,
249}
250
251impl Default for MockServerBuilder {
252    fn default() -> Self {
253        Self::new()
254    }
255}
256
257impl MockServerBuilder {
258    /// Create a new builder.
259    #[must_use]
260    pub fn new() -> Self {
261        Self {
262            name: "mock-server".to_string(),
263            version: "1.0.0".to_string(),
264            tools: Vec::new(),
265            resources: Vec::new(),
266            prompts: Vec::new(),
267        }
268    }
269
270    /// Set the server name.
271    pub fn name(mut self, name: impl Into<String>) -> Self {
272        self.name = name.into();
273        self
274    }
275
276    /// Set the server version.
277    pub fn version(mut self, version: impl Into<String>) -> Self {
278        self.version = version.into();
279        self
280    }
281
282    /// Add a mock tool.
283    #[must_use]
284    pub fn tool(mut self, tool: MockTool) -> Self {
285        self.tools.push(tool);
286        self
287    }
288
289    /// Add multiple mock tools.
290    pub fn tools(mut self, tools: impl IntoIterator<Item = MockTool>) -> Self {
291        self.tools.extend(tools);
292        self
293    }
294
295    /// Add a mock resource.
296    #[must_use]
297    pub fn resource(mut self, resource: MockResource) -> Self {
298        self.resources.push(resource);
299        self
300    }
301
302    /// Add a mock prompt.
303    #[must_use]
304    pub fn prompt(mut self, prompt: MockPrompt) -> Self {
305        self.prompts.push(prompt);
306        self
307    }
308
309    /// Build the mock server.
310    #[must_use]
311    pub fn build(self) -> MockServer {
312        let tools: HashMap<String, MockTool> = self
313            .tools
314            .into_iter()
315            .map(|t| (t.name.clone(), t))
316            .collect();
317
318        let resources: HashMap<String, MockResource> = self
319            .resources
320            .into_iter()
321            .map(|r| (r.uri.clone(), r))
322            .collect();
323
324        let prompts: HashMap<String, MockPrompt> = self
325            .prompts
326            .into_iter()
327            .map(|p| (p.name.clone(), p))
328            .collect();
329
330        MockServer {
331            name: self.name,
332            version: self.version,
333            tools,
334            resources,
335            prompts,
336        }
337    }
338}
339
340/// A mock MCP server for testing.
341///
342/// The mock server implements all handler traits and can be used
343/// with `MemoryTransport` for testing.
344pub struct MockServer {
345    name: String,
346    version: String,
347    tools: HashMap<String, MockTool>,
348    resources: HashMap<String, MockResource>,
349    prompts: HashMap<String, MockPrompt>,
350}
351
352impl MockServer {
353    /// Create a new builder for constructing a [`MockServer`].
354    ///
355    /// # Example
356    ///
357    /// ```rust
358    /// use mcpkit_testing::MockServer;
359    ///
360    /// let server = MockServer::builder()
361    ///     .name("test-server")
362    ///     .version("1.0.0")
363    ///     .build();
364    /// ```
365    #[must_use]
366    pub fn builder() -> MockServerBuilder {
367        MockServerBuilder::new()
368    }
369
370    /// Create a new builder for constructing a [`MockServer`].
371    ///
372    /// This is an alias for [`MockServer::builder()`] for backwards compatibility.
373    #[must_use]
374    #[deprecated(since = "0.2.6", note = "Use `MockServer::builder()` instead")]
375    pub fn new() -> MockServerBuilder {
376        MockServerBuilder::new()
377    }
378
379    /// Get the server name.
380    #[must_use]
381    pub fn name(&self) -> &str {
382        &self.name
383    }
384
385    /// Get the server version.
386    #[must_use]
387    pub fn version(&self) -> &str {
388        &self.version
389    }
390}
391
392impl ServerHandler for MockServer {
393    fn server_info(&self) -> ServerInfo {
394        ServerInfo::new(&self.name, &self.version)
395    }
396
397    fn capabilities(&self) -> ServerCapabilities {
398        let mut caps = ServerCapabilities::new();
399        if !self.tools.is_empty() {
400            caps = caps.with_tools();
401        }
402        if !self.resources.is_empty() {
403            caps = caps.with_resources();
404        }
405        if !self.prompts.is_empty() {
406            caps = caps.with_prompts();
407        }
408        caps
409    }
410}
411
412impl ToolHandler for MockServer {
413    fn list_tools(
414        &self,
415        _ctx: &Context,
416    ) -> impl Future<Output = Result<Vec<Tool>, McpError>> + Send {
417        let tools: Vec<Tool> = self.tools.values().map(MockTool::to_tool).collect();
418        async move { Ok(tools) }
419    }
420
421    fn call_tool(
422        &self,
423        name: &str,
424        args: Value,
425        _ctx: &Context,
426    ) -> impl Future<Output = Result<ToolOutput, McpError>> + Send {
427        let result = if let Some(tool) = self.tools.get(name) {
428            tool.call(args)
429        } else {
430            Err(McpError::method_not_found_with_suggestions(
431                name,
432                self.tools.keys().cloned().collect(),
433            ))
434        };
435        async move { result }
436    }
437}
438
439impl ResourceHandler for MockServer {
440    fn list_resources(
441        &self,
442        _ctx: &Context,
443    ) -> impl Future<Output = Result<Vec<Resource>, McpError>> + Send {
444        let resources: Vec<Resource> = self
445            .resources
446            .values()
447            .map(MockResource::to_resource)
448            .collect();
449        async move { Ok(resources) }
450    }
451
452    fn read_resource(
453        &self,
454        uri: &str,
455        _ctx: &Context,
456    ) -> impl Future<Output = Result<Vec<ResourceContents>, McpError>> + Send {
457        let result = if let Some(resource) = self.resources.get(uri) {
458            Ok(vec![resource.to_contents()])
459        } else {
460            Err(McpError::resource_not_found(uri))
461        };
462        async move { result }
463    }
464}
465
466impl PromptHandler for MockServer {
467    fn list_prompts(
468        &self,
469        _ctx: &Context,
470    ) -> impl Future<Output = Result<Vec<Prompt>, McpError>> + Send {
471        let prompts: Vec<Prompt> = self.prompts.values().map(MockPrompt::to_prompt).collect();
472        async move { Ok(prompts) }
473    }
474
475    fn get_prompt(
476        &self,
477        name: &str,
478        _args: Option<serde_json::Map<String, Value>>,
479        _ctx: &Context,
480    ) -> impl Future<Output = Result<GetPromptResult, McpError>> + Send {
481        let result = if let Some(prompt) = self.prompts.get(name) {
482            Ok(GetPromptResult {
483                description: prompt.description.clone(),
484                messages: vec![PromptMessage {
485                    role: mcpkit_core::types::Role::User,
486                    content: Content::text(&prompt.template),
487                }],
488            })
489        } else {
490            Err(McpError::method_not_found_with_suggestions(
491                name,
492                self.prompts.keys().cloned().collect(),
493            ))
494        };
495        async move { result }
496    }
497}
498
499#[cfg(test)]
500mod tests {
501    use super::*;
502
503    #[test]
504    fn test_mock_tool_text() -> Result<(), Box<dyn std::error::Error>> {
505        let tool = MockTool::new("greet").returns_text("Hello!");
506        let result = tool.call(serde_json::json!({}))?;
507        match result {
508            ToolOutput::Success(r) => {
509                assert!(!r.is_error());
510            }
511            ToolOutput::RecoverableError { .. } => panic!("Expected success"),
512        }
513        Ok(())
514    }
515
516    #[test]
517    fn test_mock_tool_error() -> Result<(), Box<dyn std::error::Error>> {
518        let tool = MockTool::new("fail").returns_error("Something went wrong");
519        let result = tool.call(serde_json::json!({}))?;
520        match result {
521            ToolOutput::RecoverableError { message, .. } => {
522                assert!(message.contains("went wrong"));
523            }
524            ToolOutput::Success(_) => panic!("Expected error"),
525        }
526        Ok(())
527    }
528
529    #[test]
530    fn test_mock_tool_dynamic() -> Result<(), Box<dyn std::error::Error>> {
531        let tool = MockTool::new("add").handler(|args| {
532            let a = args
533                .get("a")
534                .and_then(serde_json::Value::as_f64)
535                .unwrap_or(0.0);
536            let b = args
537                .get("b")
538                .and_then(serde_json::Value::as_f64)
539                .unwrap_or(0.0);
540            Ok(ToolOutput::text(format!("{}", a + b)))
541        });
542
543        let result = tool.call(serde_json::json!({"a": 1, "b": 2}))?;
544        match result {
545            ToolOutput::Success(r) => {
546                if let Content::Text(tc) = &r.content[0] {
547                    assert_eq!(tc.text, "3");
548                }
549            }
550            ToolOutput::RecoverableError { .. } => panic!("Expected success"),
551        }
552        Ok(())
553    }
554
555    #[test]
556    fn test_mock_server_builder() {
557        let server = MockServer::builder()
558            .name("test-server")
559            .version("2.0.0")
560            .tool(MockTool::new("test").returns_text("ok"))
561            .resource(MockResource::new("test://resource", "Test Resource").content("Test content"))
562            .build();
563
564        assert_eq!(server.name(), "test-server");
565        assert_eq!(server.version(), "2.0.0");
566
567        let caps = server.capabilities();
568        assert!(caps.has_tools());
569        assert!(caps.has_resources());
570        assert!(!caps.has_prompts());
571    }
572}
mcpkit_testing/mock.rs

mcpkit_testing/
mock.rs
