aimo-cli 0.0.1

AiMo Network router node Rust implementation
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
82
83
84
85
86
87
88
89
90

use std::{collections::HashMap, sync::Arc, time::Duration};

use serde_json::Value;
use tokio::time::sleep;

use crate::router::{Request, Response, Router, local::LocalRouter};

async fn echo_delay_service(request: Request, delay_ms: u64) -> serde_json::Result<Value> {
    sleep(Duration::from_micros(delay_ms)).await;

    println!(
        "Serving request {}\nbody:\n{}\nafter sleeping for {}ms",
        request.request_id, request.payload, delay_ms
    );
    serde_json::from_str::<Value>(&request.payload)
}

#[tokio::test]
async fn test_local_router() {
    // tracing_subscriber::fmt::init();

    let router = Arc::new(LocalRouter::new());

    let router_clone = router.clone();
    let jh = tokio::spawn(async move {
        router_clone.run().await;
    });

    let router_clone = router.clone();
    let srv_jh = tokio::spawn(async move {
        let mut connection = router_clone
            .register_service("test_service_id".to_string())
            .await
            .unwrap();

        loop {
            let request = connection.recv().await.unwrap();

            assert_eq!(request.service_id, "test_service_id");
            let tx = connection.tx.clone();
            tokio::spawn(async move {
                if let Ok(body) = echo_delay_service(request.clone(), 0).await {
                    tx.send(Response {
                        request_id: request.request_id,
                        status_code: 200,
                        content_type: "json".to_string(),
                        payload: body.to_string(),
                        headers: HashMap::new(),
                        is_stream_chunk: false,
                        stream_done: false,
                    })
                    .await
                    .unwrap();
                    println!("service: Response result sent");
                }
            });
        }
    });

    let router_clone = router.clone();
    tokio::spawn(async move {
        let mut reciever = router_clone
            .route_request(Request {
                sender_id: "sender".to_string(),
                request_id: "request111".to_string(),
                service_id: "test_service_id".to_string(),
                endpoint: None,
                request_type: "test".to_string(),
                payload: "{\"ping\":\"pong\"}".to_string(),
                headers: HashMap::new(),
                payload_encrypted: false,
                signature: None,
            })
            .await
            // This should not fail
            .unwrap();

        println!("Response received");

        let response = reciever.recv().await.unwrap();
        assert_eq!(response.request_id, "request111");
        assert_eq!(response.payload, "{\"ping\":\"pong\"}");

        // Done. Exit service threads.
        srv_jh.abort();
        jh.abort();
    })
    .await
    .unwrap();
}