Add integration tests for datagod-trace-agent miniagent connections

Author: Lewis-E

crates/datadog-trace-agent/tests/common/helpers.rs

@@ -0,0 +1,97 @@
+// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+//! Helper functions for integration tests
+
+use hyper::{Request, Response};
+use hyper_util::rt::TokioIo;
+use libdd_common::hyper_migration;
+use libdd_trace_utils::test_utils::create_test_json_span;
+use std::time::{Duration, UNIX_EPOCH};
+use tokio::time::timeout;
+
+/// Create a simple test trace payload as msgpack bytes
+pub fn create_test_trace_payload() -> Vec<u8> {
+    let start = UNIX_EPOCH.elapsed().unwrap().as_nanos() as i64;
+    let json_span = create_test_json_span(11, 222, 0, start, false);
+    rmp_serde::to_vec(&vec![vec![json_span]]).expect("Failed to serialize test trace")
+}
+
+/// Send an HTTP request over TCP and return the response
+pub async fn send_tcp_request(
+    port: u16,
+    uri: &str,
+    method: &str,
+    body: Option<Vec<u8>>,
+) -> Result<Response<hyper::body::Incoming>, Box<dyn std::error::Error>> {
+    let stream = timeout(
+        Duration::from_secs(2),
+        tokio::net::TcpStream::connect(format!("127.0.0.1:{}", port)),
+    )
+    .await??;
+
+    let io = TokioIo::new(stream);
+    let (mut sender, conn) = hyper::client::conn::http1::handshake(io).await?;
+
+    tokio::spawn(async move {
+        let _ = conn.await;
+    });
+
+    let mut request_builder = Request::builder()
+        .uri(uri)
+        .method(method)
+        .header("Content-Type", "application/msgpack");
+
+    let response = if let Some(body_data) = body {
+        let body_len = body_data.len();
+        request_builder = request_builder.header("Content-Length", body_len.to_string());
+        let request = request_builder.body(hyper_migration::Body::from(body_data))?;
+        timeout(Duration::from_secs(2), sender.send_request(request)).await??
+    } else {
+        let request = request_builder.body(hyper_migration::Body::empty())?;
+        timeout(Duration::from_secs(2), sender.send_request(request)).await??
+    };
+
+    Ok(response)
+}
+
+#[cfg(windows)]
+/// Send an HTTP request over named pipe and return the response
+pub async fn send_named_pipe_request(
+    pipe_name: &str,
+    uri: &str,
+    method: &str,
+    body: Option<Vec<u8>>,
+) -> Result<Response<hyper::body::Incoming>, Box<dyn std::error::Error>> {
+    use tokio::net::windows::named_pipe::ClientOptions;
+
+    let client = timeout(
+        Duration::from_secs(2),
+        ClientOptions::new().open(pipe_name),
+    )
+    .await??;
+
+    let io = TokioIo::new(client);
+    let (mut sender, conn) = hyper::client::conn::http1::handshake(io).await?;
+
+    tokio::spawn(async move {
+        let _ = conn.await;
+    });
+
+    let mut request_builder = Request::builder()
+        .uri(uri)
+        .method(method)
+        .header("Content-Type", "application/msgpack");
+
+    let response = if let Some(body_data) = body {
+        let body_len = body_data.len();
+        request_builder = request_builder.header("Content-Length", body_len.to_string());
+        let request = request_builder.body(hyper_migration::Body::from(body_data))?;
+        timeout(Duration::from_secs(2), sender.send_request(request)).await??
+    } else {
+        let request = request_builder.body(hyper_migration::Body::empty())?;
+        timeout(Duration::from_secs(2), sender.send_request(request)).await??
+    };
+
+    Ok(response)
+}

crates/datadog-trace-agent/tests/common/mocks.rs

@@ -0,0 +1,118 @@
+// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+//! Mock implementations of trace agent components for testing
+
+use datadog_trace_agent::{
+    config::Config,
+    env_verifier::EnvVerifier,
+    stats_flusher::StatsFlusher,
+    stats_processor::StatsProcessor,
+    trace_flusher::TraceFlusher,
+    trace_processor::TraceProcessor,
+};
+use libdd_common::hyper_migration;
+use libdd_trace_protobuf::pb;
+use libdd_trace_utils::trace_utils::{self, MiniAgentMetadata, SendData};
+use std::sync::Arc;
+use tokio::sync::mpsc::{Receiver, Sender};
+
+/// Mock trace processor that returns 200 OK for all requests
+#[allow(dead_code)]
+pub struct MockTraceProcessor;
+
+#[async_trait::async_trait]
+impl TraceProcessor for MockTraceProcessor {
+    async fn process_traces(
+        &self,
+        _config: Arc<Config>,
+        _req: hyper_migration::HttpRequest,
+        _trace_tx: Sender<SendData>,
+        _mini_agent_metadata: Arc<MiniAgentMetadata>,
+    ) -> Result<hyper_migration::HttpResponse, hyper::http::Error> {
+        hyper::Response::builder()
+            .status(200)
+            .body(hyper_migration::Body::from("{}"))
+    }
+}
+
+/// Mock trace flusher that consumes messages without processing
+pub struct MockTraceFlusher;
+
+#[async_trait::async_trait]
+impl TraceFlusher for MockTraceFlusher {
+    fn new(
+        _aggregator: Arc<tokio::sync::Mutex<datadog_trace_agent::aggregator::TraceAggregator>>,
+        _config: Arc<Config>,
+    ) -> Self {
+        MockTraceFlusher
+    }
+
+    async fn start_trace_flusher(&self, mut trace_rx: Receiver<SendData>) {
+        // Consume messages from the channel without processing them
+        while let Some(_trace) = trace_rx.recv().await {
+            // Just discard the trace - we're not testing the flusher
+        }
+    }
+
+    async fn send(&self, _traces: Vec<SendData>) -> Option<Vec<SendData>> {
+        None
+    }
+
+    async fn flush(&self, _failed_traces: Option<Vec<SendData>>) -> Option<Vec<SendData>> {
+        None
+    }
+}
+
+/// Mock stats processor that returns 200 OK for all requests
+pub struct MockStatsProcessor;
+
+#[async_trait::async_trait]
+impl StatsProcessor for MockStatsProcessor {
+    async fn process_stats(
+        &self,
+        _config: Arc<Config>,
+        _req: hyper_migration::HttpRequest,
+        _stats_tx: Sender<pb::ClientStatsPayload>,
+    ) -> Result<hyper_migration::HttpResponse, hyper::http::Error> {
+        hyper::Response::builder()
+            .status(200)
+            .body(hyper_migration::Body::from("{}"))
+    }
+}
+
+/// Mock stats flusher that consumes messages without processing
+pub struct MockStatsFlusher;
+
+#[async_trait::async_trait]
+impl StatsFlusher for MockStatsFlusher {
+    async fn start_stats_flusher(
+        &self,
+        _config: Arc<Config>,
+        mut stats_rx: Receiver<pb::ClientStatsPayload>,
+    ) {
+        // Consume messages from the channel without processing them
+        while let Some(_stats) = stats_rx.recv().await {
+            // Just discard the stats - we're not testing the flusher
+        }
+    }
+
+    async fn flush_stats(&self, _config: Arc<Config>, _traces: Vec<pb::ClientStatsPayload>) {
+        // Do nothing
+    }
+}
+
+/// Mock environment verifier that returns default metadata
+pub struct MockEnvVerifier;
+
+#[async_trait::async_trait]
+impl EnvVerifier for MockEnvVerifier {
+    async fn verify_environment(
+        &self,
+        _timeout_ms: u64,
+        _env_type: &trace_utils::EnvironmentType,
+        _os: &str,
+    ) -> MiniAgentMetadata {
+        MiniAgentMetadata::default()
+    }
+}

crates/datadog-trace-agent/tests/common/mod.rs

@@ -0,0 +1,7 @@
+// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+//! Common test utilities, mocks, and helpers for integration tests
+
+pub mod helpers;
+pub mod mocks;