otel-arrow-dfe-engine 0.61.0

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

//! End-to-end tests for the [`extension_capabilities!`] macro.
//!
//! Each test defines a hand-rolled test capability (the `#[capability]`
//! proc macro cannot expand inside the engine crate itself -- proc
//! macros can't target paths in their host crate), invokes one arm of
//! the macro, exercises the resulting `register_*` fn pointers against
//! a real `CapabilityRegistry`, then resolves bindings and consumes the
//! capability.

use super::registry::{
    Capabilities, CapabilityRegistry, ConsumedTracker, LocalCapabilityEntry, SharedCapabilityEntry,
    resolve_bindings,
};
use super::{ExtensionCapability, KNOWN_CAPABILITIES, KnownCapability};
use crate::capability::factory::{LocalInstanceFactory, SharedInstanceFactory};
use otel_arrow_dfe_config::{CapabilityId, ExtensionId};
use std::any::{Any, TypeId};
use std::collections::{HashMap, HashSet};

// -- Test capability (mirrors what `#[capability]` generates) ---------

trait MacroTestCapLocal {
    fn value(&self) -> &str;
}
trait MacroTestCapShared: Send + Sync {
    fn value(&self) -> &str;
}

struct MacroTestCap;
impl super::private::Sealed for MacroTestCap {}
impl ExtensionCapability for MacroTestCap {
    const NAME: &'static str = "macro_test_cap";
    type Local = dyn MacroTestCapLocal;
    type Shared = dyn MacroTestCapShared;
    fn wrap_shared_as_local(shared: Box<Self::Shared>) -> Box<Self::Local> {
        struct Adapter(Box<dyn MacroTestCapShared>);
        impl MacroTestCapLocal for Adapter {
            fn value(&self) -> &str {
                self.0.value()
            }
        }
        Box::new(Adapter(shared))
    }
}

// Register the test capability so resolve_bindings sees it as known.
#[allow(unsafe_code)]
#[linkme::distributed_slice(KNOWN_CAPABILITIES)]
#[linkme(crate = linkme)]
static _MACRO_TEST_CAP: KnownCapability = KnownCapability {
    name: "macro_test_cap",
    description: "Test capability for extension_capabilities! macro tests",
    type_id: || TypeId::of::<MacroTestCap>(),
};

// Casters mirror what `#[capability]` emits. The macro calls
// `<MacroTestCap>::shared_entry::<E>` and `local_entry::<E>`.
impl MacroTestCap {
    fn shared_entry<E>(ext_id: ExtensionId, factory: SharedInstanceFactory) -> SharedCapabilityEntry
    where
        E: MacroTestCapShared + 'static,
    {
        let produce = move || -> Box<dyn Any + Send> {
            let erased = factory.produce();
            let concrete: Box<E> = erased.downcast().expect("instance factory");
            let shared: Box<dyn MacroTestCapShared> = concrete;
            Box::new(shared) as Box<dyn Any + Send>
        };
        let adapt_as_local: fn(Box<dyn Any + Send>) -> Box<dyn Any> = |erased| {
            let shared: Box<Box<dyn MacroTestCapShared>> = erased.downcast().expect("envelope");
            let boxed_local = <MacroTestCap as ExtensionCapability>::wrap_shared_as_local(*shared);
            Box::new(boxed_local) as Box<dyn Any>
        };
        SharedCapabilityEntry::new(ext_id, produce, adapt_as_local)
    }

    fn local_entry<E>(ext_id: ExtensionId, factory: LocalInstanceFactory) -> LocalCapabilityEntry
    where
        E: MacroTestCapLocal + 'static,
    {
        let produce = move || -> Box<dyn Any> {
            let erased = factory.produce();
            let concrete: Box<E> = erased.downcast().expect("instance factory");
            let local: Box<dyn MacroTestCapLocal> = concrete;
            Box::new(local) as Box<dyn Any>
        };
        LocalCapabilityEntry::new(ext_id, produce)
    }
}

// -- Test extension impls ---------------------------------------------

#[derive(Clone)]
struct Shared(&'static str);
impl MacroTestCapShared for Shared {
    fn value(&self) -> &str {
        self.0
    }
}

#[derive(Clone)]
struct Local(&'static str);
impl MacroTestCapLocal for Local {
    fn value(&self) -> &str {
        self.0
    }
}

fn shared_factory(val: &'static str) -> SharedInstanceFactory {
    SharedInstanceFactory::new(move || Box::new(Shared(val)) as Box<dyn Any + Send>)
}
fn local_factory(val: &'static str) -> LocalInstanceFactory {
    LocalInstanceFactory::new(move || Box::new(Local(val)) as Box<dyn Any>)
}

fn bindings() -> HashMap<CapabilityId, ExtensionId> {
    let mut m = HashMap::new();
    let _ = m.insert("macro_test_cap".into(), "ext".into());
    m
}
fn known_exts() -> HashSet<ExtensionId> {
    let mut s = HashSet::new();
    let _ = s.insert("ext".into());
    s
}

fn resolve(registry: &CapabilityRegistry) -> Capabilities {
    let mut tracker = ConsumedTracker::new();
    resolve_bindings(&bindings(), registry, &known_exts(), &mut tracker).expect("resolve_bindings")
}

// -- Tests ------------------------------------------------------------

#[test]
fn macro_shared_only_form() {
    let ec = extension_capabilities!(shared: Shared => [MacroTestCap]);
    assert_eq!(ec.shared, &["macro_test_cap"]);
    assert!(ec.local.is_empty());

    let mut registry = CapabilityRegistry::new();
    (ec.register_shared)("ext".into(), shared_factory("s-only"), &mut registry)
        .expect("register_shared");
    // register_local must be a no-op even when called with a local factory.
    (ec.register_local)("ext".into(), local_factory("unused"), &mut registry)
        .expect("register_local no-op");

    // Native shared claim works.
    let caps = resolve(&registry);
    assert_eq!(
        caps.require_shared::<MacroTestCap>().unwrap().value(),
        "s-only"
    );

    // Fresh resolve: local consumers are served by the
    // SharedAsLocal fallback, returning the same value.
    let caps = resolve(&registry);
    assert_eq!(
        caps.require_local::<MacroTestCap>().unwrap().value(),
        "s-only"
    );

    // Per-binding one-shot: claiming the local execution model via
    // the SharedAsLocal fallback consumes the binding for the
    // native shared execution model too.
    let caps = resolve(&registry);
    let _ = caps.require_local::<MacroTestCap>().unwrap();
    assert!(matches!(
        caps.require_shared::<MacroTestCap>(),
        Err(crate::capability::registry::Error::CapabilityAlreadyConsumed { .. })
    ));
}

#[test]
fn macro_local_only_form() {
    let ec = extension_capabilities!(local: Local => [MacroTestCap]);
    assert!(ec.shared.is_empty());
    assert_eq!(ec.local, &["macro_test_cap"]);

    let mut registry = CapabilityRegistry::new();
    (ec.register_local)("ext".into(), local_factory("l-only"), &mut registry)
        .expect("register_local");

    let caps = resolve(&registry);
    assert_eq!(
        caps.require_local::<MacroTestCap>().unwrap().value(),
        "l-only"
    );
    assert!(caps.require_shared::<MacroTestCap>().is_err());
}

#[test]
fn macro_dual_form() {
    let ec = extension_capabilities!(
        (shared: Shared, local: Local) => [MacroTestCap]
    );
    assert_eq!(ec.shared, &["macro_test_cap"]);
    assert_eq!(ec.local, &["macro_test_cap"]);

    let mut registry = CapabilityRegistry::new();
    (ec.register_shared)("ext".into(), shared_factory("s-dual"), &mut registry)
        .expect("register_shared");
    (ec.register_local)("ext".into(), local_factory("l-dual"), &mut registry)
        .expect("register_local");

    // Native dual: local and shared are distinct registrations
    // (different concrete types), but the per-binding one-shot
    // contract applies uniformly -- claiming either execution model
    // on a node invalidates the other on the same node. To observe
    // both concrete values we resolve twice, claiming a different
    // side from each independent `Capabilities`.
    let caps_shared = resolve(&registry);
    assert_eq!(
        caps_shared
            .require_shared::<MacroTestCap>()
            .unwrap()
            .value(),
        "s-dual"
    );
    assert!(matches!(
        caps_shared.require_local::<MacroTestCap>(),
        Err(crate::capability::registry::Error::CapabilityAlreadyConsumed { .. })
    ));

    let caps_local = resolve(&registry);
    assert_eq!(
        caps_local.require_local::<MacroTestCap>().unwrap().value(),
        "l-dual"
    );
    assert!(matches!(
        caps_local.require_shared::<MacroTestCap>(),
        Err(crate::capability::registry::Error::CapabilityAlreadyConsumed { .. })
    ));

    // Each execution model is still one-shot in isolation.
    let caps = resolve(&registry);
    let _ = caps.require_shared::<MacroTestCap>().unwrap();
    assert!(matches!(
        caps.require_shared::<MacroTestCap>(),
        Err(crate::capability::registry::Error::CapabilityAlreadyConsumed { .. })
    ));
}