1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

use std::convert::Infallible;

use futures_util::{future::BoxFuture, SinkExt, StreamExt};
use tower::Service;

use crate::{
    box_error,
    client::transport::SocketChannels,
    common::{socket::SocketMessage, transport::mock::MockReceiver},
    response::BoxResponse,
    server::{socket::SocketHandler, MakeRoutes},
};

use super::Transport;

/// Mock transport for the server.
pub struct Mock {
    rx: MockReceiver,
}

impl Mock {
    /// Create a new mock transport.
    pub fn new(rx: MockReceiver) -> Self {
        Self { rx }
    }
}

impl Transport for Mock {
    type Error = Infallible;

    fn serve<S>(mut self, mk_routes: S) -> BoxFuture<'static, Result<(), Self::Error>>
    where
        S: MakeRoutes,
    {
        let mut svc = mk_routes
            .into_make_service()
            .call(())
            .into_inner()
            .unwrap()
            .unwrap();

        Box::pin(async move {
            while let Some((req, sender)) = self.rx.inner.recv().await {
                let fut = Service::call(&mut svc, req);

                tokio::spawn(async move {
                    let mut resp = fut.await.expect("cant fail");

                    if let Some(socket_handler) = resp.extensions_mut().remove::<SocketHandler>() {
                        let (client_tx, server_rx) =
                            futures_channel::mpsc::unbounded::<SocketMessage>();
                        let (server_tx, client_rx) = futures_channel::mpsc::unbounded();

                        {
                            let client_socket_chans = SocketChannels::new(
                                Box::pin(client_tx.sink_map_err(box_error)),
                                Box::pin(client_rx.map(Ok)),
                            );

                            let mut resp = BoxResponse::empty();
                            resp.extensions_mut().insert(client_socket_chans);

                            sender.send(resp).expect("sender dropped");
                        }

                        let fut = (socket_handler.inner)(
                            Box::pin(server_rx.map(Ok)),
                            Box::pin(server_tx.sink_map_err(box_error)),
                        );

                        tokio::spawn(fut);
                    } else {
                        sender.send(resp).expect("sender dropped");
                    }
                });
            }

            Ok(())
        })
    }
}