Refactor

src/tools/docs/docs.rs

@@ -11,28 +11,25 @@ use reqwest::Client;
 use serde_json::{json, Value};
 use tokio::sync::Mutex;
 
-#[cfg(test)]
-mod tests;
-
 // Cache for documentation lookups to avoid repeated requests
 #[derive(Clone)]
-struct DocCache {
+pub struct DocCache {
     cache: Arc<Mutex<std::collections::HashMap<String, String>>>,
 }
 
 impl DocCache {
-    fn new() -> Self {
+    pub fn new() -> Self {
         Self {
             cache: Arc::new(Mutex::new(std::collections::HashMap::new())),
         }
     }
 
-    async fn get(&self, key: &str) -> Option<String> {
+    pub async fn get(&self, key: &str) -> Option<String> {
         let cache = self.cache.lock().await;
         cache.get(key).cloned()
     }
 
-    async fn set(&self, key: String, value: String) {
+    pub async fn set(&self, key: String, value: String) {
         let mut cache = self.cache.lock().await;
         cache.insert(key, value);
     }

src/tools/docs/mod.rs

@@ -0,0 +1,4 @@
+pub mod docs;
+mod tests;
+
+pub use docs::DocRouter;

src/tools/docs/tests.rs

@@ -1,6 +1,5 @@
-use super::*;
+use crate::tools::{DocRouter, DocCache};
 use mcp_core::{Content, ToolError};
-use mcp_server::Router;
 use serde_json::json;
 
 #[tokio::test]

src/tools/mod.rs

@@ -1,3 +1,4 @@
-mod docs;
+pub mod docs;
 
-pub use docs::DocRouter;
+pub use docs::DocRouter;
+pub use docs::docs::DocCache;

src/transport/http_sse_server/http_sse_server.rs

@@ -6,14 +6,14 @@ use axum::{
     routing::get,
     Router,
 };
-
-#[cfg(test)]
-mod tests;
-use futures::{stream::Stream, StreamExt, TryStreamExt};
+use futures::{Stream, StreamExt, TryStreamExt};
 use mcp_server::{ByteTransport, Server};
 use std::collections::HashMap;
 use tokio_util::codec::FramedRead;
 
+#[cfg(test)]
+// Tests in ../tests.rs
+
 use anyhow::Result;
 use mcp_server::router::RouterService;
 use crate::{transport::jsonrpc_frame_codec::JsonRpcFrameCodec, tools::DocRouter};

src/transport/http_sse_server/mod.rs

@@ -0,0 +1,4 @@
+mod http_sse_server;
+mod tests;
+
+pub use http_sse_server::*;

src/transport/http_sse_server/tests.rs

@@ -1,18 +1,11 @@
-use super::*;
-use axum::{
-    body::Body,
-    http::{Method, Request},
-};
-use tokio::sync::RwLock;
+use crate::transport::http_sse_server::App;
 // Comment out tower imports for now, as we'll handle router testing differently
 // use tower::Service; 
 // use tower::util::ServiceExt;
 
 // Helper function to create an App with an empty state
 fn create_test_app() -> App {
-    App {
-        txs: Arc::new(RwLock::new(HashMap::new())),
-    }
+    App::new()
 }
 
 #[tokio::test]

src/transport/jsonrpc_frame_codec/jsonrpc_frame_codec.rs

@@ -3,8 +3,6 @@ use tokio_util::codec::Decoder;
 #[derive(Default)]
 pub struct JsonRpcFrameCodec;
 
-#[cfg(test)]
-mod tests;
 impl Decoder for JsonRpcFrameCodec {
     type Item = tokio_util::bytes::Bytes;
     type Error = tokio::io::Error;

src/transport/jsonrpc_frame_codec/mod.rs

@@ -0,0 +1,4 @@
+mod jsonrpc_frame_codec;
+mod tests;
+
+pub use jsonrpc_frame_codec::JsonRpcFrameCodec;

src/transport/jsonrpc_frame_codec/tests.rs

@@ -1,5 +1,6 @@
-use super::*;
+use crate::transport::jsonrpc_frame_codec::JsonRpcFrameCodec;
 use tokio_util::bytes::BytesMut;
+use tokio_util::codec::Decoder;
 
 #[test]
 fn test_decode_single_line() {

tests/integration_tests.rs

@@ -1,5 +1,4 @@
 use cratedocs_mcp::{tools::DocRouter, transport::jsonrpc_frame_codec::JsonRpcFrameCodec};
-use mcp_core::Tool;
 use mcp_server::Router;
 use serde_json::{json, Value};
 use tokio_util::codec::Decoder;