use core::convert::Infallible;
use crate::{address::SocketAddress, client::EstablishedClientConnection, stream::Socket};
use rama_core::{
Service,
error::BoxError,
extensions::{Extensions, ExtensionsRef},
rt::Executor,
};
use tokio::io::{AsyncRead, AsyncWrite, DuplexStream, duplex};
#[derive(Debug, Clone)]
pub struct MockConnectorService<S> {
create_server: S,
max_buffer_size: usize,
executor: Option<Executor>,
}
impl<S> MockConnectorService<S> {
pub fn new(create_server: S) -> Self {
Self {
create_server,
max_buffer_size: 1024,
executor: None,
}
}
rama_utils::macros::generate_set_and_with! {
pub fn max_buffer_size(mut self, size: usize) -> Self {
self.max_buffer_size = size;
self
}
}
rama_utils::macros::generate_set_and_with! {
pub fn executor(mut self, executor: Option<Executor>) -> Self {
self.executor = executor;
self
}
}
}
impl<S, Input, Server> Service<Input> for MockConnectorService<S>
where
S: Fn() -> Server + Send + Sync + 'static,
Server: Service<MockSocket, Error: Into<BoxError>>,
Input: Send + 'static,
{
type Error = Infallible;
type Output = EstablishedClientConnection<MockSocket, Input>;
async fn serve(&self, input: Input) -> Result<Self::Output, Self::Error> {
let (client, server) = duplex(self.max_buffer_size);
let client_socket = MockSocket::new(client);
let server_socket = MockSocket::new(server);
let server = (self.create_server)();
self.executor
.clone()
.unwrap_or_default()
.into_spawn_task(async move {
#[expect(
clippy::panic,
reason = "test helper: failures must surface loudly to invalidate the test run"
)]
if let Err(err) = server.serve(server_socket).await {
panic!("created mock server failed: {}", err.into())
}
});
Ok(EstablishedClientConnection {
input,
conn: client_socket,
})
}
}
#[derive(Debug)]
pub struct MockSocket {
stream: DuplexStream,
extensions: Extensions,
}
impl MockSocket {
#[must_use]
pub fn new(stream: DuplexStream) -> Self {
Self {
stream,
extensions: Extensions::new(),
}
}
}
impl ExtensionsRef for MockSocket {
fn extensions(&self) -> &Extensions {
&self.extensions
}
}
#[warn(clippy::missing_trait_methods)]
impl AsyncRead for MockSocket {
fn poll_read(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::pin::Pin::new(&mut self.stream).poll_read(cx, buf)
}
}
#[warn(clippy::missing_trait_methods)]
impl AsyncWrite for MockSocket {
fn poll_write(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
buf: &[u8],
) -> std::task::Poll<std::io::Result<usize>> {
std::pin::Pin::new(&mut self.stream).poll_write(cx, buf)
}
fn poll_flush(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::pin::Pin::new(&mut self.stream).poll_flush(cx)
}
fn poll_shutdown(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<()>> {
std::pin::Pin::new(&mut self.stream).poll_shutdown(cx)
}
fn is_write_vectored(&self) -> bool {
self.stream.is_write_vectored()
}
fn poll_write_vectored(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
bufs: &[std::io::IoSlice<'_>],
) -> std::task::Poll<Result<usize, std::io::Error>> {
std::pin::Pin::new(&mut self.stream).poll_write_vectored(cx, bufs)
}
}
impl Socket for MockSocket {
fn local_addr(&self) -> std::io::Result<SocketAddress> {
Ok(SocketAddress::local_ipv4(0))
}
fn peer_addr(&self) -> std::io::Result<SocketAddress> {
Ok(SocketAddress::local_ipv4(0))
}
}