Refactor

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() {