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

mcpkit_testing/
mock.rs
