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.
354    #[must_use]
355    pub fn builder() -> MockServerBuilder {
356        MockServerBuilder::new()
357    }
358
359    /// Create a simple mock server.
360    #[must_use]
361    pub fn new() -> MockServerBuilder {
362        MockServerBuilder::new()
363    }
364
365    /// Get the server name.
366    #[must_use]
367    pub fn name(&self) -> &str {
368        &self.name
369    }
370
371    /// Get the server version.
372    #[must_use]
373    pub fn version(&self) -> &str {
374        &self.version
375    }
376}
377
378impl ServerHandler for MockServer {
379    fn server_info(&self) -> ServerInfo {
380        ServerInfo::new(&self.name, &self.version)
381    }
382
383    fn capabilities(&self) -> ServerCapabilities {
384        let mut caps = ServerCapabilities::new();
385        if !self.tools.is_empty() {
386            caps = caps.with_tools();
387        }
388        if !self.resources.is_empty() {
389            caps = caps.with_resources();
390        }
391        if !self.prompts.is_empty() {
392            caps = caps.with_prompts();
393        }
394        caps
395    }
396}
397
398impl ToolHandler for MockServer {
399    fn list_tools(
400        &self,
401        _ctx: &Context,
402    ) -> impl Future<Output = Result<Vec<Tool>, McpError>> + Send {
403        let tools: Vec<Tool> = self.tools.values().map(MockTool::to_tool).collect();
404        async move { Ok(tools) }
405    }
406
407    fn call_tool(
408        &self,
409        name: &str,
410        args: Value,
411        _ctx: &Context,
412    ) -> impl Future<Output = Result<ToolOutput, McpError>> + Send {
413        let result = if let Some(tool) = self.tools.get(name) {
414            tool.call(args)
415        } else {
416            Err(McpError::method_not_found_with_suggestions(
417                name,
418                self.tools.keys().cloned().collect(),
419            ))
420        };
421        async move { result }
422    }
423}
424
425impl ResourceHandler for MockServer {
426    fn list_resources(
427        &self,
428        _ctx: &Context,
429    ) -> impl Future<Output = Result<Vec<Resource>, McpError>> + Send {
430        let resources: Vec<Resource> = self
431            .resources
432            .values()
433            .map(MockResource::to_resource)
434            .collect();
435        async move { Ok(resources) }
436    }
437
438    fn read_resource(
439        &self,
440        uri: &str,
441        _ctx: &Context,
442    ) -> impl Future<Output = Result<Vec<ResourceContents>, McpError>> + Send {
443        let result = if let Some(resource) = self.resources.get(uri) {
444            Ok(vec![resource.to_contents()])
445        } else {
446            Err(McpError::resource_not_found(uri))
447        };
448        async move { result }
449    }
450}
451
452impl PromptHandler for MockServer {
453    fn list_prompts(
454        &self,
455        _ctx: &Context,
456    ) -> impl Future<Output = Result<Vec<Prompt>, McpError>> + Send {
457        let prompts: Vec<Prompt> = self.prompts.values().map(MockPrompt::to_prompt).collect();
458        async move { Ok(prompts) }
459    }
460
461    fn get_prompt(
462        &self,
463        name: &str,
464        _args: Option<serde_json::Map<String, Value>>,
465        _ctx: &Context,
466    ) -> impl Future<Output = Result<GetPromptResult, McpError>> + Send {
467        let result = if let Some(prompt) = self.prompts.get(name) {
468            Ok(GetPromptResult {
469                description: prompt.description.clone(),
470                messages: vec![PromptMessage {
471                    role: mcpkit_core::types::Role::User,
472                    content: Content::text(&prompt.template),
473                }],
474            })
475        } else {
476            Err(McpError::method_not_found_with_suggestions(
477                name,
478                self.prompts.keys().cloned().collect(),
479            ))
480        };
481        async move { result }
482    }
483}
484
485#[cfg(test)]
486mod tests {
487    use super::*;
488
489    #[test]
490    fn test_mock_tool_text() {
491        let tool = MockTool::new("greet").returns_text("Hello!");
492        let result = tool.call(serde_json::json!({})).unwrap();
493        match result {
494            ToolOutput::Success(r) => {
495                assert!(!r.is_error());
496            }
497            ToolOutput::RecoverableError { .. } => panic!("Expected success"),
498        }
499    }
500
501    #[test]
502    fn test_mock_tool_error() {
503        let tool = MockTool::new("fail").returns_error("Something went wrong");
504        let result = tool.call(serde_json::json!({})).unwrap();
505        match result {
506            ToolOutput::RecoverableError { message, .. } => {
507                assert!(message.contains("went wrong"));
508            }
509            ToolOutput::Success(_) => panic!("Expected error"),
510        }
511    }
512
513    #[test]
514    fn test_mock_tool_dynamic() {
515        let tool = MockTool::new("add").handler(|args| {
516            let a = args
517                .get("a")
518                .and_then(serde_json::Value::as_f64)
519                .unwrap_or(0.0);
520            let b = args
521                .get("b")
522                .and_then(serde_json::Value::as_f64)
523                .unwrap_or(0.0);
524            Ok(ToolOutput::text(format!("{}", a + b)))
525        });
526
527        let result = tool.call(serde_json::json!({"a": 1, "b": 2})).unwrap();
528        match result {
529            ToolOutput::Success(r) => {
530                if let Content::Text(tc) = &r.content[0] {
531                    assert_eq!(tc.text, "3");
532                }
533            }
534            ToolOutput::RecoverableError { .. } => panic!("Expected success"),
535        }
536    }
537
538    #[test]
539    fn test_mock_server_builder() {
540        let server = MockServer::new()
541            .name("test-server")
542            .version("2.0.0")
543            .tool(MockTool::new("test").returns_text("ok"))
544            .resource(MockResource::new("test://resource", "Test Resource").content("Test content"))
545            .build();
546
547        assert_eq!(server.name(), "test-server");
548        assert_eq!(server.version(), "2.0.0");
549
550        let caps = server.capabilities();
551        assert!(caps.has_tools());
552        assert!(caps.has_resources());
553        assert!(!caps.has_prompts());
554    }
555}
mcpkit_testing/mock.rs

mcpkit_testing/
mock.rs
