otel-arrow-dfe-engine 0.61.0

// Copyright The OpenTelemetry Authors
// SPDX-License-Identifier: Apache-2.0

//! Common testing utilities for engine components.
//!
//! This module provides shared testing constructs used across tests for receivers,
//! processors, and exporters. It includes:
//!
//! - Shared types like `TestMsg` for passing data through the pipeline in tests
//! - Counter mechanisms for tracking control message processing
//! - Utilities for setting up single-threaded async test runtimes
//! - Channel creation helpers for connecting components
//!
//! The specialized testing utilities for receivers, processors, and exporters are in their respective
//! submodules.

use crate::Interests;
use crate::attributes::{ExtensionScopeAttributeSet, PipelineAttributeSet};
use crate::context::{ControllerContext, ExtensionContext, PipelineContext};
use crate::control::NodeControlMsg;
use crate::runtime_services::PipelineRuntimeServices;
use otel_arrow_dfe_channel::mpsc;
use otel_arrow_dfe_config::engine::{
    ResolvedOtelDataflowSpec, ResolvedPipelineConfig, ResolvedPipelineRole,
};
use otel_arrow_dfe_config::node::NodeKind;
use otel_arrow_dfe_config::pipeline::PipelineConfig;
use otel_arrow_dfe_config::policy::{DistributionTier, Policies};
use otel_arrow_dfe_telemetry::registry::TelemetryRegistryHandle;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use tokio::runtime::Builder;
use tokio::task::LocalSet;

pub mod capability;
#[cfg(any(test, feature = "test-utils"))]
pub mod dst;
pub mod exporter;
pub mod liveness;
pub mod node;
pub mod processor;
pub mod receiver;

pub use node::{test_node, test_nodes};

/// Creates fresh pipeline runtime services for one logical test pipeline.
///
/// Tests that construct multiple effect handlers for the same pipeline should clone the returned
/// value and inject one clone into each handler.
///
/// # Panics
///
/// Panics when the test binary links invalid pdata codec registrations. Such registrations are a
/// test assembly error; tests for registry validation should construct the registry directly.
#[cfg(any(test, feature = "test-utils"))]
#[must_use]
pub fn test_pipeline_runtime_services() -> PipelineRuntimeServices {
    create_test_pipeline_runtime_services()
}

fn create_test_pipeline_runtime_services() -> PipelineRuntimeServices {
    PipelineRuntimeServices::new(Default::default())
        .expect("test binary must link valid pdata codec registrations")
}

/// Create a minimal [`PipelineContext`] suitable for unit tests that
/// need to register metrics or construct engine objects.
///
/// Returns both the context and the [`TelemetryRegistryHandle`] so
/// callers can inspect registered metrics.
#[must_use]
pub fn test_pipeline_ctx() -> (PipelineContext, TelemetryRegistryHandle) {
    test_pipeline_ctx_with_interests(Interests::empty())
}

/// Create a minimal [`PipelineContext`] with explicit node interests for tests.
#[must_use]
pub fn test_pipeline_ctx_with_interests(
    interests: Interests,
) -> (PipelineContext, TelemetryRegistryHandle) {
    test_pipeline_ctx_with_interests_and_duration_distribution(interests, DistributionTier::Normal)
}

/// Create a minimal [`PipelineContext`] with explicit node interests and local duration tier.
#[must_use]
pub fn test_pipeline_ctx_with_interests_and_duration_distribution(
    interests: Interests,
    duration_distribution: DistributionTier,
) -> (PipelineContext, TelemetryRegistryHandle) {
    let registry = TelemetryRegistryHandle::new();
    let controller = ControllerContext::new(registry.clone());
    let mut ctx = controller
        .pipeline_context_with("test_grp".into(), "test_pipeline".into(), 0, 1, 0)
        .with_node_context(
            "test_node".into(),
            "urn:test:processor:example".into(),
            NodeKind::Processor,
            HashMap::new(),
        );
    ctx.set_node_interests(interests);
    ctx.set_node_duration_distribution(duration_distribution);
    (ctx, registry)
}

/// Compiles and installs bindings for one test pipeline.
///
/// Ignores engine and group policies.
/// Compile the full engine configuration to test inheritance or multiple pipelines.
///
/// # Errors
///
/// Returns configuration or declaration errors.
pub fn install_test_context_bindings<PData: 'static + Clone + std::fmt::Debug>(
    pipeline_ctx: &mut PipelineContext,
    factory: &crate::PipelineFactory<PData>,
    pipeline: PipelineConfig,
) -> Result<(), crate::error::Error> {
    let resolved = ResolvedOtelDataflowSpec {
        engine: Default::default(),
        pipelines: vec![ResolvedPipelineConfig {
            pipeline_group_id: pipeline_ctx.pipeline_group_id(),
            pipeline_id: pipeline_ctx.pipeline_id(),
            policies: Policies::resolve(pipeline.policies()),
            pipeline,
            topic_scope: otel_arrow_dfe_config::topic::TopicScope::GLOBAL,
            role: ResolvedPipelineRole::Regular,
        }],
    };
    pipeline_ctx
        .set_compiled_context_bindings(factory.compile_initial_context(&resolved)?.bindings);
    Ok(())
}

/// Create a minimal [`ExtensionContext`] suitable for unit tests of the
/// extension subsystem. Returns the context and the underlying registry.
#[must_use]
pub fn test_extension_ctx() -> (ExtensionContext, TelemetryRegistryHandle) {
    let registry = TelemetryRegistryHandle::new();
    let controller = ControllerContext::new(registry.clone());
    let scope = ExtensionScopeAttributeSet::pipeline(PipelineAttributeSet {
        pipeline_group_id: "test_group".into(),
        pipeline_id: "test_pipeline".into(),
        ..PipelineAttributeSet::default()
    });
    (ExtensionContext::new(controller, scope), registry)
}

/// Create a minimal extension effect handler for tests that run an extension directly.
#[cfg(any(test, feature = "test-utils"))]
#[must_use]
pub fn test_extension_effect_handler(
    name: otel_arrow_dfe_config::ExtensionId,
) -> crate::extension::EffectHandler {
    let (tx, _rx) = flume::bounded(1);
    crate::extension::EffectHandler::new(
        name,
        otel_arrow_dfe_telemetry::reporter::MetricsReporter::new(tx),
        None,
    )
}

/// A test message type used in component tests.
#[derive(Debug, PartialEq, Clone)]
pub struct TestMsg(pub String);

impl crate::ReceivedAtNode for TestMsg {
    fn received_at_node(&mut self, _node_id: usize, _node_interests: Interests) {}
}

impl crate::processor::FlowMetricHook for TestMsg {}

impl TestMsg {
    /// Creates a new test message with the given content.
    pub fn new<S: Into<String>>(content: S) -> Self {
        TestMsg(content.into())
    }
}

/// Set of counters for tracking the number of control messages processed.
#[derive(Clone)]
pub struct CtrlMsgCounters {
    timer_tick_count: Arc<AtomicUsize>,
    message_count: Arc<AtomicUsize>,
    config_count: Arc<AtomicUsize>,
    shutdown_count: Arc<AtomicUsize>,
}

impl CtrlMsgCounters {
    /// Creates a new set of counters with all counts initialized to zero.
    #[must_use]
    pub fn new() -> Self {
        CtrlMsgCounters {
            timer_tick_count: Arc::new(AtomicUsize::new(0)),
            message_count: Arc::new(AtomicUsize::new(0)),
            config_count: Arc::new(AtomicUsize::new(0)),
            shutdown_count: Arc::new(AtomicUsize::new(0)),
        }
    }

    /// Handles incoming control messages and increments the appropriate counter.
    pub fn update_with<PData>(&self, msg: &NodeControlMsg<PData>) {
        match msg {
            NodeControlMsg::TimerTick { .. } => self.increment_timer_tick(),
            NodeControlMsg::Config { .. } => self.increment_config(),
            NodeControlMsg::Shutdown { .. } => self.increment_shutdown(),
            _ => {}
        }
    }

    /// Increments the timer tick count.
    pub fn increment_timer_tick(&self) {
        _ = self
            .timer_tick_count
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    /// Increments the message count.
    pub fn increment_message(&self) {
        _ = self
            .message_count
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    /// Increments the config count.
    pub fn increment_config(&self) {
        _ = self
            .config_count
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    /// Increments the shutdown count.
    pub fn increment_shutdown(&self) {
        _ = self
            .shutdown_count
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    /// Gets the current timer tick count.
    #[must_use]
    pub fn get_timer_tick_count(&self) -> usize {
        self.timer_tick_count
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Gets the current message count.
    #[must_use]
    pub fn get_message_count(&self) -> usize {
        self.message_count
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Gets the current config count.
    #[must_use]
    pub fn get_config_count(&self) -> usize {
        self.config_count.load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Gets the current shutdown count.
    #[must_use]
    pub fn get_shutdown_count(&self) -> usize {
        self.shutdown_count
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Asserts that the current counters match the expected values.
    pub fn assert(
        &self,
        timer_tick_count: usize,
        message_count: usize,
        config_count: usize,
        shutdown_count: usize,
    ) {
        assert_eq!(
            self.get_timer_tick_count(),
            timer_tick_count,
            "Timer tick count mismatch"
        );
        assert_eq!(
            self.get_message_count(),
            message_count,
            "Message count mismatch"
        );
        assert_eq!(
            self.get_config_count(),
            config_count,
            "Config count mismatch"
        );
        assert_eq!(
            self.get_shutdown_count(),
            shutdown_count,
            "Shutdown count mismatch"
        );
    }
}

impl Default for CtrlMsgCounters {
    fn default() -> Self {
        Self::new()
    }
}

/// Creates a single-threaded runtime with a local task set for testing components.
#[must_use]
pub fn setup_test_runtime() -> (tokio::runtime::Runtime, LocalSet) {
    // Check if we're already inside a Tokio runtime
    if tokio::runtime::Handle::try_current().is_ok() {
        panic!(
            "TestRuntime used inside an existing async runtime: change #[tokio::test] to #[test]."
        );
    }

    let rt = Builder::new_current_thread()
        .enable_all()
        .build()
        .expect("Failed to create new runtime");
    let local_tasks = LocalSet::new();
    (rt, local_tasks)
}

/// Helper to create `!Send` MPSC channels with a specific capacity.
///
/// This function creates a sender-receiver pair with the given capacity.
#[must_use]
pub fn create_not_send_channel<T>(capacity: usize) -> (mpsc::Sender<T>, mpsc::Receiver<T>) {
    mpsc::Channel::new(capacity)
}