rmon/rmon-agent/tests/agent_test.rs

use anyhow::Result;
use tokio::net::{TcpStream, UdpSocket};
use tokio_util::codec::Framed;
use futures::{StreamExt, SinkExt};
use rmon_common::protocol::{ClientMessage, AgentMessage, BincodeCodec, PROTOCOL_VERSION};
use std::time::Duration;
use tokio::time::sleep;
use bytes::BytesMut;
use tokio_util::codec::{Decoder, Encoder};
use std::io;
use rmon_agent::config::Config;
use rmon_agent::run_agent;
use std::path::PathBuf;
use std::fs;

/// Client side codec
pub struct ClientCodec;

impl Decoder for ClientCodec {
    type Item = AgentMessage;
    type Error = io::Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        let mut codec = BincodeCodec::<AgentMessage>::default();
        codec.decode(src)
    }
}

impl Encoder<ClientMessage> for ClientCodec {
    type Error = io::Error;

    fn encode(&mut self, item: ClientMessage, dst: &mut BytesMut) -> Result<(), Self::Error> {
        let mut codec = BincodeCodec::<ClientMessage>::default();
        codec.encode(item, dst)
    }
}

async fn setup_agent(port: u16, udp_port: u16, config_path: PathBuf, data_dir: &str) -> Result<()> {
    let config_str = format!(r#"
[storage]
data_dir = "{}"
default_mode = "circular"
default_window = "1h"

[[sources]]
type = "udp"
name = "test_source"
bind = "127.0.0.1:{}"
[sources.timestamp]
format = "none"
[[sources.signals]]
id = "sig1"
label = "Sig 1"
unit = "V"
scale = 1.0
offset = 0.0
byte_offset = 0
byte_type = "f32_le"
path = ["test"]
"#, data_dir, udp_port);

    if let Some(parent) = config_path.parent() {
        fs::create_dir_all(parent)?;
    }
    fs::write(&config_path, config_str.clone())?;
    let config: Config = toml::from_str(&config_str)?;

    tokio::spawn(async move {
        let _ = run_agent(config, &format!("127.0.0.1:{}", port), config_path).await;
    });

    sleep(Duration::from_millis(300)).await;
    Ok(())
}

#[tokio::test]
async fn test_agent_full_flow() -> Result<()> {
    tokio::time::timeout(Duration::from_secs(20), async {
        let port = 7892;
        let udp_port = 9001;
        let data_dir = "/tmp/rmon-test-data-full";
        let config_path = PathBuf::from("/tmp/rmon-test/config_full.toml");
        let _ = fs::remove_dir_all(data_dir);
        
        setup_agent(port, udp_port, config_path, data_dir).await?;

        let stream = TcpStream::connect(format!("127.0.0.1:{}", port)).await?;
        let mut framed = Framed::new(stream, ClientCodec);

        framed.send(ClientMessage::Handshake {
            protocol_version: PROTOCOL_VERSION,
            client_version: "0.1.0".to_string(),
        }).await?;

        let ack = framed.next().await.ok_or_else(|| anyhow::anyhow!("Disconnected"))??;
        if let AgentMessage::HandshakeAck { accepted, .. } = ack {
            assert!(accepted);
        } else {
            panic!("Expected HandshakeAck, got {:?}", ack);
        }

        framed.send(ClientMessage::Subscribe {
            ids: vec!["udp://test_source/sig1".to_string()],
            from_ns: None,
        }).await?;

        let udp_socket = UdpSocket::bind("127.0.0.1:0").await?;
        let val: f32 = 42.0;
        udp_socket.send_to(&val.to_le_bytes(), format!("127.0.0.1:{}", udp_port)).await?;

        let msg = framed.next().await.ok_or_else(|| anyhow::anyhow!("Disconnected"))??;
        if let AgentMessage::DataBatch { batches } = msg {
            assert_eq!(batches[0].id, "udp://test_source/sig1");
            assert_eq!(batches[0].samples[0].value, 42.0);
        } else {
            panic!("Expected DataBatch, got {:?}", msg);
        }

        // Check continuous storage
        sleep(Duration::from_millis(200)).await;
        let signal_dir = format!("{}/test_source/udp___test_source_sig1", data_dir);
        let data_files = fs::read_dir(&signal_dir).map_err(|e| anyhow::anyhow!("Failed to read {}: {}", signal_dir, e))?;
        assert!(data_files.count() > 0);

        Ok::<(), anyhow::Error>(())
    }).await.map_err(|_| anyhow::anyhow!("Test timed out"))?
}

#[tokio::test]
async fn test_agent_config_remote() -> Result<()> {
    tokio::time::timeout(Duration::from_secs(20), async {
        let port = 7894;
        let udp_port = 9002;
        let data_dir = "/tmp/rmon-test-data-remote";
        let config_path = PathBuf::from("/tmp/rmon-test/config_remote.toml");
        let _ = fs::remove_dir_all(data_dir);
        
        setup_agent(port, udp_port, config_path.clone(), data_dir).await?;

        let stream = TcpStream::connect(format!("127.0.0.1:{}", port)).await?;
        let mut framed = Framed::new(stream, ClientCodec);

        framed.send(ClientMessage::Handshake {
            protocol_version: PROTOCOL_VERSION,
            client_version: "0.1.0".to_string(),
        }).await?;
        let _ = framed.next().await.unwrap()?;

        // 1. Get Config
        framed.send(ClientMessage::GetConfig).await?;
        let msg = framed.next().await.ok_or_else(|| anyhow::anyhow!("Disconnected"))??;
        if let AgentMessage::ConfigReport { config_toml } = msg {
            assert!(config_toml.contains("test_source"));
        } else {
            panic!("Expected ConfigReport, got {:?}", msg);
        }

        // 2. Update Config
        let new_config = format!(r#"
[storage]
data_dir = "{}"
default_mode = "circular"
default_window = "1h"

[[sources]]
type = "shell"
name = "new_source"
command = "echo \"val: 100\""
poll_interval = "100ms"
[sources.timestamp]
format = "none"
[[sources.signals]]
id = "sig_new"
label = "New Sig"
unit = "X"
scale = 1.0
offset = 0.0
regex = "val: ([0-9]+)"
path = ["new"]
"#, data_dir);
        framed.send(ClientMessage::UpdateConfig { config_toml: new_config.to_string() }).await?;
        
        // Wait for hot-reload
        sleep(Duration::from_millis(2000)).await;

        // 3. Check if signals changed
        framed.send(ClientMessage::ListSignals).await?;
        let msg = tokio::time::timeout(Duration::from_secs(5), framed.next()).await
            .map_err(|_| anyhow::anyhow!("Timeout waiting for SignalList"))?
            .ok_or_else(|| anyhow::anyhow!("Disconnected"))??;
            
        if let AgentMessage::SignalList { signals } = msg {
            assert!(signals.iter().any(|s| s.id == "shell://new_source/sig_new"), "New signal not found in {:?}", signals);
        } else {
            panic!("Expected SignalList, got {:?}", msg);
        }

        Ok::<(), anyhow::Error>(())
    }).await.map_err(|_| anyhow::anyhow!("Test timed out"))?
}