Skip to main content

mockforge_ws/
protocol_server.rs

1//! Unified protocol server implementation for the WebSocket mock server.
2
3use async_trait::async_trait;
4use mockforge_core::protocol_abstraction::Protocol;
5use mockforge_core::protocol_server::MockProtocolServer;
6
7use crate::LatencyProfile;
8
9/// A `MockProtocolServer` wrapper around the WebSocket server startup.
10///
11/// Wraps [`crate::start_with_latency_and_host`] with shutdown-signal integration.
12pub struct WsMockServer {
13    port: u16,
14    host: String,
15    latency: Option<LatencyProfile>,
16}
17
18impl WsMockServer {
19    /// Create a new `WsMockServer` with the given configuration.
20    pub fn new(port: u16, host: String, latency: Option<LatencyProfile>) -> Self {
21        Self {
22            port,
23            host,
24            latency,
25        }
26    }
27}
28
29#[async_trait]
30impl MockProtocolServer for WsMockServer {
31    fn protocol(&self) -> Protocol {
32        Protocol::WebSocket
33    }
34
35    async fn start(
36        &self,
37        mut shutdown: tokio::sync::watch::Receiver<()>,
38    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
39        let port = self.port;
40        let host = self.host.clone();
41        let latency = self.latency.clone();
42
43        tokio::select! {
44            result = crate::start_with_latency_and_host(port, &host, latency) => {
45                result.map_err(|e| -> Box<dyn std::error::Error + Send + Sync> {
46                    Box::new(std::io::Error::other(e.to_string()))
47                })
48            }
49            _ = shutdown.changed() => {
50                tracing::info!("Shutting down WebSocket server on port {}", port);
51                Ok(())
52            }
53        }
54    }
55
56    fn port(&self) -> u16 {
57        self.port
58    }
59
60    fn description(&self) -> String {
61        format!("WebSocket server on {}:{}", self.host, self.port)
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn test_ws_mock_server_protocol() {
71        let server = WsMockServer::new(3001, "0.0.0.0".to_string(), None);
72        assert_eq!(server.protocol(), Protocol::WebSocket);
73    }
74
75    #[test]
76    fn test_ws_mock_server_port() {
77        let server = WsMockServer::new(3001, "0.0.0.0".to_string(), None);
78        assert_eq!(server.port(), 3001);
79    }
80
81    #[test]
82    fn test_ws_mock_server_description() {
83        let server = WsMockServer::new(3001, "127.0.0.1".to_string(), None);
84        assert_eq!(server.description(), "WebSocket server on 127.0.0.1:3001");
85    }
86}