add basic test

own_mcp/src/mcp/chat.rs

@@ -1,16 +1,15 @@
 use crate::mcp::MCPClient;
-use ollama_rs::Ollama;
 use ollama_rs::error::OllamaError;
-use ollama_rs::generation::chat::ChatMessage;
 use ollama_rs::generation::chat::request::ChatMessageRequest;
+use ollama_rs::generation::chat::ChatMessage;
 use ollama_rs::generation::tools::ToolInfo;
-use rmcp::ServiceError;
+use ollama_rs::Ollama;
 use rmcp::model::{CallToolRequestParams, CallToolResult};
+use rmcp::ServiceError;
 use std::collections::HashMap;
-use std::ops::Index;
 use thiserror::Error;
 
-#[derive(Debug, Copy, Clone, Default)]
+#[derive(Debug, Copy, Clone, Default, PartialEq)]
 pub enum ToolPermission {
     Allowed,
     #[default]
@@ -138,6 +137,13 @@ impl AgentChat {
     /// returns:
     /// Ok(tuple) => tuple of the server name and the tool or resource name
     /// Err(error) => ChatError::FunctionParseError(name)
+    ///
+    /// # Examples
+    ///
+    /// ```ignore
+    /// let result = AgentChat::parse_tool_name("example::get_foo".to_string());
+    /// assert_eq!(result.unwrap(), ("example".to_string(), "get_foo".to_string()))
+    /// ```
     fn parse_tool_name(name: String) -> Result<(String, String), ChatError> {
         let (mcp_server_name, tool_name) = name
             .split_once("::")
@@ -158,11 +164,10 @@ impl AgentChat {
     /// * `name`: must be in the format "mcp_server_name::tool_or_resource_name".
     ///
     /// returns: Option<&RestrictedTool>
-    pub fn get_tool<'a>(
-        &self,
-        mpc_server_data: &'a MCPServerData,
+    pub fn get_tool(
+        mpc_server_data: &MCPServerData,
         name: String,
-    ) -> Option<&'a RestrictedTool> {
+    ) -> Option<&RestrictedTool> {
         mpc_server_data
             .translated_tools
             .iter()
@@ -248,11 +253,6 @@ impl AgentChat {
     ///
     /// returns: Result<Option<CallToolResult>, ChatError>
     /// if permission was denied: Ok(None)
-    /// # Examples
-    ///
-    /// ```
-    ///
-    /// ```
     async fn call_tool<C: Future<Output = PermissionAnswer>>(
         &self,
         tool_call: &ollama_rs::generation::tools::ToolCall,
@@ -273,8 +273,7 @@ impl AgentChat {
         let mcp_server_data = self
             .get_mcp_server_by_name(mcp_server_name.clone())
             .ok_or(ChatError::ServiceNotFoundError(mcp_server_name.clone()))?;
-        let restricted_tool = self
-            .get_tool(&mcp_server_data, full_unparsed_tool_name.clone())
+        let restricted_tool = AgentChat::get_tool(&mcp_server_data, full_unparsed_tool_name.clone())
             .ok_or(ChatError::ToolNotFoundError(
                 full_unparsed_tool_name.clone(),
             ))?;
@@ -314,3 +313,48 @@ impl AgentChat {
         Ok(Some(mcp_server_data.client.call_tool(request_params).await?))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::mcp::chat::RestrictedTool;
+    use crate::mcp::get_client;
+    use rmcp::model::Implementation;
+    use std::collections::HashMap;
+
+    const MCP_TEST_SERVER_URL: &str = "https://mcpplaygroundonline.com/mcp-echo-server";
+
+    #[tokio::test]
+    async fn test_permission_setting() {
+        let ollama = Ollama::default();
+
+        println!(
+            "Test requires ollama running on {} and connection to {MCP_TEST_SERVER_URL}",
+            ollama.url()
+        );
+
+        let model_name = "lfm2.5-thinking:1.2b".into();
+
+        ollama.pull_model(model_name, false).await.unwrap();
+
+        let mcp_clients = HashMap::from([(
+            "test".to_string(),
+            get_client(
+                MCP_TEST_SERVER_URL,
+                None::<String>,
+                Implementation::new("test", env!("CARGO_PKG_VERSION")),
+            ).await.unwrap(),
+        )]);
+
+        let mut chat = AgentChat::new(ollama, "lfm2.5-thinking:1.2b".to_string(), mcp_clients, String::new()).await.unwrap();
+
+        let tools: Vec<RestrictedTool> = chat.get_all_tools().cloned().collect();
+        assert_eq!(tools.len(), 1);
+
+        assert_eq!(tools[0].permission, ToolPermission::Ask);
+        chat.set_permission("test::echo".to_string(), ToolPermission::Allowed).unwrap();
+
+        let tools: Vec<RestrictedTool> = chat.get_all_tools().cloned().collect();
+        assert_eq!(tools[0].permission, ToolPermission::Allowed);
+    }
+}