eagle/tests/client.rs

/*
Eagle - A simple library for RPC in Rust
Copyright (c) 2024 KodiCraft

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
use eagle::Protocol;
use env_logger::{Builder, Env};
use std::sync::Arc;
use std::sync::Once;
use tokio::sync::{
    mpsc::{self, Receiver, Sender},
    Notify,
};

static INIT: Once = Once::new();
pub fn init_logger() {
    INIT.call_once(|| {
        let env = Env::default()
            .filter_or("RUST_LOG", "info")
            .write_style_or("LOG_STYLE", "always");

        Builder::from_env(env).format_timestamp_nanos().init();
    });
}

#[derive(Protocol)]
enum TestProtocol {
    Addition((i32, i32), i32),
    SomeKindOfQuestion(String, i32),
    ThisRespondsWithAString(i32, String),
    Void((), ()),
}

#[tokio::test]
async fn client() {
    init_logger();
    let (qtx, qrx) = mpsc::channel(16);
    let (atx, arx) = mpsc::channel(16);
    let ready_notify = Arc::new(Notify::new());
    let client = TestProtocolClient::new(qtx, arx, None, ready_notify.clone());
    ready_notify.notify_one();
    let server = tokio::spawn(server_loop(qrx, atx));
    let result = client.addition(2, 5).await.unwrap();
    assert_eq!(result, 7);
    let result = client
        .some_kind_of_question("Hello, world!".to_string())
        .await
        .unwrap();
    assert_eq!(result, "Hello, world!".len() as i32);
    let result = client.this_responds_with_a_string(42).await.unwrap();
    assert_eq!(result, "The number is 42");
    client.void().await.unwrap();
    server.abort();
}

async fn server_loop(
    mut qrx: Receiver<(u64, __TestProtocolQuestion)>,
    atx: Sender<(u64, __TestProtocolAnswer)>,
) {
    loop {
        if let Some((nonce, q)) = qrx.recv().await {
            match q {
                __TestProtocolQuestion::addition((a, b)) => {
                    atx.send((nonce, __TestProtocolAnswer::addition(a + b)))
                        .await
                        .unwrap();
                }
                __TestProtocolQuestion::some_kind_of_question(s) => {
                    atx.send((
                        nonce,
                        __TestProtocolAnswer::some_kind_of_question(s.len() as i32),
                    ))
                    .await
                    .unwrap();
                }
                __TestProtocolQuestion::this_responds_with_a_string(i) => {
                    atx.send((
                        nonce,
                        __TestProtocolAnswer::this_responds_with_a_string(format!(
                            "The number is {}",
                            i
                        )),
                    ))
                    .await
                    .unwrap();
                }
                __TestProtocolQuestion::void(()) => {
                    println!("Received void question");
                    atx.send((nonce, __TestProtocolAnswer::void(())))
                        .await
                        .unwrap();
                }
            }
        }
    }
}

#[tokio::test]
async fn heavy_async() {
    init_logger();
    let (qtx, qrx) = mpsc::channel(16);
    let (atx, arx) = mpsc::channel(16);
    let ready_notify = Arc::new(Notify::new());
    let client = TestProtocolClient::new(qtx, arx, None, ready_notify.clone());
    ready_notify.notify_one();
    let server = tokio::spawn(server_loop(qrx, atx));
    let mut tasks = Vec::new();
    for i in 0..100 {
        let client = client.clone();
        tasks.push(tokio::spawn(async move {
            tokio::time::sleep(tokio::time::Duration::from_millis(
                rand::random::<u64>() % 100,
            ))
            .await;
            let result = client.addition(i, i).await.unwrap();
            assert_eq!(result, i + i);
        }));
    }
    for task in tasks {
        task.await.unwrap();
    }
    server.abort();
}