observer_rust_lib/

lib.rs

// SPDX-License-Identifier: MIT

use anyhow::{anyhow, Result};
use base64::engine::general_purpose::STANDARD;
use base64::Engine;
use observer_rust::{TelemetryEntry as RustTelemetryEntry, TelemetryValue as RustTelemetryValue};
use regex::Regex;
use serde::Serialize;
use serde_json::Number;
use std::cell::RefCell;
use std::collections::{BTreeSet, HashMap};
use std::fmt::Debug;
use std::panic::{catch_unwind, panic_any, resume_unwind, AssertUnwindSafe};

pub const TEST_EXIT_PASS: i32 = 0;

pub type TestOutcome = observer_rust::TestOutcome;
pub type TestFunction = fn(&mut TestContext);

#[derive(Debug, Clone)]
pub struct TestRegistration {
    title: String,
    canonical_name: String,
    target: String,
    function: TestFunction,
    suite_path: Vec<String>,
}

impl TestRegistration {
    pub fn title(&self) -> &str {
        &self.title
    }

    pub fn canonical_name(&self) -> &str {
        &self.canonical_name
    }

    pub fn target(&self) -> &str {
        &self.target
    }

    pub fn suite_path(&self) -> &[String] {
        &self.suite_path
    }
}

#[derive(Debug, Clone, Default)]
pub struct TestOptions<'a> {
    pub id: Option<&'a str>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RegistryError {
    EmptyCanonicalName,
    EmptyTarget,
    DuplicateCanonicalName(String),
    DuplicateTarget(String),
}

impl std::fmt::Display for RegistryError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::EmptyCanonicalName => {
                write!(f, "registered test canonical name must not be empty")
            }
            Self::EmptyTarget => write!(f, "registered test target must not be empty"),
            Self::DuplicateCanonicalName(value) => {
                write!(f, "duplicate canonical test name `{value}`")
            }
            Self::DuplicateTarget(value) => write!(f, "duplicate test target `{value}`"),
        }
    }
}

impl std::error::Error for RegistryError {}

#[derive(Debug)]
pub struct ExpectationFailure {
    message: String,
}

impl ExpectationFailure {
    pub fn message(&self) -> &str {
        &self.message
    }
}

pub struct TestContext {
    inner: observer_rust::TestContext,
}

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

impl TestContext {
    pub fn new() -> Self {
        Self {
            inner: observer_rust::TestContext::new(),
        }
    }

    pub fn stdout<S>(&mut self, value: S)
    where
        S: AsRef<[u8]>,
    {
        self.inner.write_out(value.as_ref());
    }

    pub fn stderr<S>(&mut self, value: S)
    where
        S: AsRef<[u8]>,
    {
        self.inner.write_err(value.as_ref());
    }

    pub fn set_exit(&mut self, exit: i32) {
        self.inner.set_exit(exit);
    }

    pub fn fail(&mut self, message: &str) {
        self.inner.fail_now(message);
    }

    pub fn observe(&mut self) -> Observe<'_> {
        Observe {
            inner: &mut self.inner,
        }
    }

    pub fn as_raw_mut(&mut self) -> &mut observer_rust::TestContext {
        &mut self.inner
    }

    pub fn finish(self) -> TestOutcome {
        self.inner.finish()
    }
}

pub struct Observe<'a> {
    inner: &'a mut observer_rust::TestContext,
}

impl Observe<'_> {
    pub fn metric(&mut self, name: &str, value: f64) -> bool {
        self.inner.emit_metric(name, value)
    }

    pub fn vector(&mut self, name: &str, values: &[f64]) -> bool {
        self.inner.emit_vector(name, values)
    }

    pub fn tag(&mut self, name: &str, value: &str) -> bool {
        self.inner.emit_tag(name, value)
    }
}

pub trait Truthy {
    fn is_truthy(&self) -> bool;
}

impl Truthy for bool {
    fn is_truthy(&self) -> bool {
        *self
    }
}

impl<T> Truthy for Option<T> {
    fn is_truthy(&self) -> bool {
        self.is_some()
    }
}

impl<T, E> Truthy for Result<T, E> {
    fn is_truthy(&self) -> bool {
        self.is_ok()
    }
}

pub trait StrPattern {
    fn matches_input(&self, input: &str) -> bool;
    fn describe_pattern(&self) -> String;
}

impl StrPattern for &str {
    fn matches_input(&self, input: &str) -> bool {
        input.contains(self)
    }

    fn describe_pattern(&self) -> String {
        (*self).to_owned()
    }
}

impl StrPattern for String {
    fn matches_input(&self, input: &str) -> bool {
        input.contains(self)
    }

    fn describe_pattern(&self) -> String {
        self.clone()
    }
}

impl StrPattern for Regex {
    fn matches_input(&self, input: &str) -> bool {
        self.is_match(input)
    }

    fn describe_pattern(&self) -> String {
        self.as_str().to_owned()
    }
}

pub struct Expectation<T> {
    actual: T,
    inverted: bool,
}

pub fn expect<T>(actual: T) -> Expectation<T> {
    Expectation {
        actual,
        inverted: false,
    }
}

impl<T> Expectation<T> {
    pub fn not(mut self) -> Self {
        self.inverted = !self.inverted;
        self
    }

    fn apply(&self, condition: bool, message: String) {
        if self.inverted {
            if condition {
                panic_any(ExpectationFailure {
                    message: format!("not {message}"),
                });
            }
            return;
        }

        if !condition {
            panic_any(ExpectationFailure { message });
        }
    }
}

impl<T> Expectation<T>
where
    T: PartialEq + Debug,
{
    pub fn to_be(self, expected: T) {
        let message = format!("expected {:?} to be {:?}", &self.actual, &expected);
        self.apply(self.actual == expected, message);
    }

    pub fn to_equal(self, expected: T) {
        let message = format!("expected {:?} to equal {:?}", &self.actual, &expected);
        self.apply(self.actual == expected, message);
    }
}

impl<T> Expectation<T>
where
    T: Truthy + Debug,
{
    pub fn to_be_truthy(self) {
        let message = format!("expected {:?} to be truthy", &self.actual);
        self.apply(self.actual.is_truthy(), message);
    }

    pub fn to_be_falsy(self) {
        let message = format!("expected {:?} to be falsy", &self.actual);
        self.apply(!self.actual.is_truthy(), message);
    }
}

impl<T> Expectation<T>
where
    T: AsRef<str> + Debug,
{
    pub fn to_contain<S>(self, expected: S)
    where
        S: AsRef<str> + Debug,
    {
        let actual = self.actual.as_ref();
        let expected_ref = expected.as_ref();
        let message = format!("expected {:?} to contain {:?}", actual, expected_ref);
        self.apply(actual.contains(expected_ref), message);
    }

    pub fn to_match<P>(self, pattern: P)
    where
        P: StrPattern,
    {
        let actual = self.actual.as_ref();
        let description = pattern.describe_pattern();
        let message = format!("expected {:?} to match {:?}", actual, description);
        self.apply(pattern.matches_input(actual), message);
    }
}

#[derive(Default)]
struct CollectionState {
    tests: Vec<CollectedTest>,
    suite_path: Vec<String>,
}

#[derive(Clone)]
struct CollectedTest {
    title: String,
    id: Option<String>,
    function: TestFunction,
    suite_path: Vec<String>,
}

thread_local! {
    static ACTIVE_COLLECTION: RefCell<Option<CollectionState>> = RefCell::new(None);
}

fn with_collection_state<T>(f: impl FnOnce(&mut CollectionState) -> T) -> T {
    ACTIVE_COLLECTION.with(|cell| {
        let mut borrowed = cell.borrow_mut();
        let state = borrowed
            .as_mut()
            .expect("test collection is only available inside collect_tests(...)");
        f(state)
    })
}

fn derive_base_id(suite_path: &[String], title: &str) -> String {
    if suite_path.is_empty() {
        title.to_owned()
    } else {
        format!("{} :: {title}", suite_path.join(" :: "))
    }
}

fn materialize_registrations(collected: Vec<CollectedTest>) -> Vec<TestRegistration> {
    let mut counts = HashMap::<String, usize>::new();

    collected
        .into_iter()
        .map(|test| {
            let base_id = test
                .id
                .clone()
                .unwrap_or_else(|| derive_base_id(&test.suite_path, &test.title));
            let next_count = counts.get(&base_id).copied().unwrap_or(0) + 1;
            counts.insert(base_id.clone(), next_count);
            let resolved_id = if test.id.is_some() || next_count == 1 {
                base_id
            } else {
                format!("{base_id} #{next_count}")
            };

            TestRegistration {
                title: test.title,
                canonical_name: resolved_id.clone(),
                target: resolved_id,
                function: test.function,
                suite_path: test.suite_path,
            }
        })
        .collect()
}

pub fn define_test(
    title: &str,
    function: TestFunction,
    options: TestOptions<'_>,
) -> Result<TestRegistration, RegistryError> {
    let tests = materialize_registrations(vec![CollectedTest {
        title: title.to_owned(),
        id: options.id.map(str::to_owned),
        function,
        suite_path: Vec::new(),
    }]);
    Ok(sorted_validated_tests(&tests)?.remove(0))
}

pub fn collect_tests(define: impl FnOnce()) -> Result<Vec<TestRegistration>, RegistryError> {
    let previous = ACTIVE_COLLECTION.with(|cell| cell.replace(Some(CollectionState::default())));
    let define_result = catch_unwind(AssertUnwindSafe(define));
    let state = ACTIVE_COLLECTION
        .with(|cell| cell.replace(previous))
        .expect("collection state should exist");

    match define_result {
        Ok(()) => sorted_validated_tests(&materialize_registrations(state.tests)),
        Err(payload) => resume_unwind(payload),
    }
}

pub fn describe(label: &str, define: impl FnOnce()) {
    with_collection_state(|state| state.suite_path.push(label.to_owned()));
    let define_result = catch_unwind(AssertUnwindSafe(define));
    with_collection_state(|state| {
        state.suite_path.pop();
    });

    if let Err(payload) = define_result {
        resume_unwind(payload);
    }
}

pub fn register_test(title: &str, id: Option<&str>, function: TestFunction) {
    with_collection_state(|state| {
        state.tests.push(CollectedTest {
            title: title.to_owned(),
            id: id.map(str::to_owned),
            function,
            suite_path: state.suite_path.clone(),
        });
    });
}

pub fn test(title: &str, function: TestFunction) {
    register_test(title, None, function);
}

pub fn it(title: &str, function: TestFunction) {
    register_test(title, None, function);
}

pub fn sorted_validated_tests(tests: &[TestRegistration]) -> Result<Vec<TestRegistration>, RegistryError> {
    validate_registry(tests)?;
    let mut sorted = tests.to_vec();
    sorted.sort_by(|left, right| {
        left.canonical_name
            .as_bytes()
            .cmp(right.canonical_name.as_bytes())
            .then_with(|| left.target.as_bytes().cmp(right.target.as_bytes()))
    });
    Ok(sorted)
}

pub fn validate_registry(tests: &[TestRegistration]) -> Result<(), RegistryError> {
    let mut names = BTreeSet::new();
    let mut targets = BTreeSet::new();

    for test in tests {
        if test.canonical_name.trim().is_empty() {
            return Err(RegistryError::EmptyCanonicalName);
        }
        if test.target.trim().is_empty() {
            return Err(RegistryError::EmptyTarget);
        }
        if !names.insert(test.canonical_name.as_str()) {
            return Err(RegistryError::DuplicateCanonicalName(
                test.canonical_name.clone(),
            ));
        }
        if !targets.insert(test.target.as_str()) {
            return Err(RegistryError::DuplicateTarget(test.target.clone()));
        }
    }

    Ok(())
}

pub fn find_target<'a>(tests: &'a [TestRegistration], target: &str) -> Option<&'a TestRegistration> {
    tests.iter().find(|test| test.target == target)
}

pub fn run_test(test: &TestRegistration) -> TestOutcome {
    let mut ctx = TestContext::new();
    let result = catch_unwind_silent(AssertUnwindSafe(|| {
        (test.function)(&mut ctx);
    }));

    match result {
        Ok(()) => ctx.finish(),
        Err(payload) => {
            let message = panic_payload_message(&payload);
            ctx.fail(&message);
            ctx.finish()
        }
    }
}

fn catch_unwind_silent<F, R>(f: F) -> std::thread::Result<R>
where
    F: FnOnce() -> R + std::panic::UnwindSafe,
{
    let hook = std::panic::take_hook();
    std::panic::set_hook(Box::new(|_| {}));
    let result = catch_unwind(f);
    std::panic::set_hook(hook);
    result
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct HostRunArgs {
    pub target: String,
    pub timeout_ms: u32,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct ListPayload {
    pub provider: String,
    pub tests: Vec<ListEntry>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct ListEntry {
    pub name: String,
    pub target: String,
}

#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct RunPayload {
    pub provider: String,
    pub target: String,
    pub exit: i32,
    pub out_b64: String,
    pub err_b64: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub telemetry: Option<Vec<TelemetryPayloadEntry>>,
}

#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct TelemetryPayloadEntry {
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub unit: Option<String>,
    #[serde(flatten)]
    pub value: TelemetryPayloadValue,
}

#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum TelemetryPayloadValue {
    Metric { value: Number },
    Vector { values: Vec<Number> },
    Tag { value: String },
}

pub fn list_payload(provider: &str, tests: &[TestRegistration]) -> Result<ListPayload> {
    validate_provider(provider)?;
    let sorted = sorted_validated_tests(tests).map_err(|error| anyhow!(error))?;
    Ok(ListPayload {
        provider: provider.to_owned(),
        tests: sorted
            .into_iter()
            .map(|test| ListEntry {
                name: test.canonical_name,
                target: test.target,
            })
            .collect(),
    })
}

pub fn run_payload(
    provider: &str,
    tests: &[TestRegistration],
    target: &str,
    timeout_ms: u32,
) -> Result<RunPayload> {
    let _ = timeout_ms;
    validate_provider(provider)?;
    let sorted = sorted_validated_tests(tests).map_err(|error| anyhow!(error))?;
    let test = find_target(&sorted, target).ok_or_else(|| anyhow!("unknown test target `{target}`"))?;
    let outcome = run_test(test);
    Ok(RunPayload {
        provider: provider.to_owned(),
        target: target.to_owned(),
        exit: outcome.exit,
        out_b64: STANDARD.encode(&outcome.out),
        err_b64: STANDARD.encode(&outcome.err),
        telemetry: telemetry_payload_from_outcome(outcome),
    })
}

pub fn observer_host_handles_command(command: Option<&str>) -> bool {
    matches!(command, Some("list" | "run" | "observe"))
}

pub fn parse_run_args(args: &[String]) -> Result<HostRunArgs> {
    let mut target = None;
    let mut timeout_ms = 0;
    let mut index = 0;

    while index < args.len() {
        match args[index].as_str() {
            "--target" => {
                let value = args
                    .get(index + 1)
                    .ok_or_else(|| anyhow!("missing --target value"))?;
                target = Some(value.clone());
                index += 2;
            }
            "--timeout-ms" => {
                let value = args
                    .get(index + 1)
                    .ok_or_else(|| anyhow!("missing --timeout-ms value"))?;
                timeout_ms = value
                    .parse::<u32>()
                    .map_err(|_| anyhow!("invalid --timeout-ms value"))?;
                index += 2;
            }
            other => return Err(anyhow!("unknown argument {other}")),
        }
    }

    Ok(HostRunArgs {
        target: target.ok_or_else(|| anyhow!("missing --target"))?,
        timeout_ms,
    })
}

pub fn observer_host_dispatch(provider: &str, tests: &[TestRegistration], argv: &[String]) -> Result<()> {
    validate_provider(provider)?;

    let Some(command) = argv.first().map(String::as_str) else {
        return Err(anyhow!("expected subcommand list|run|observe"));
    };

    match command {
        "list" => write_json(&list_payload(provider, tests)?),
        "run" | "observe" => {
            let args = parse_run_args(&argv[1..])?;
            write_json(&run_payload(provider, tests, &args.target, args.timeout_ms)?)
        }
        other => Err(anyhow!("unknown subcommand `{other}`")),
    }
}

pub fn observer_host_main(provider: &str, tests: &[TestRegistration]) -> Result<()> {
    observer_host_main_from(provider, tests, std::env::args())
}

pub fn observer_host_main_from<I, S>(
    provider: &str,
    tests: &[TestRegistration],
    argv: I,
) -> Result<()>
where
    I: IntoIterator<Item = S>,
    S: Into<String>,
{
    let args = argv.into_iter().map(Into::into).collect::<Vec<_>>();
    observer_host_dispatch(provider, tests, args.get(1..).unwrap_or(&[]))
}

pub fn observer_host_dispatch_embedded<I, S>(
    provider: &str,
    root_command: &str,
    tests: &[TestRegistration],
    argv: I,
) -> Result<()>
where
    I: IntoIterator<Item = S>,
    S: Into<String>,
{
    if root_command.trim().is_empty() {
        return Err(anyhow!("root command must not be empty"));
    }

    let args = argv.into_iter().map(Into::into).collect::<Vec<_>>();
    if args.len() < 2 || args[1] != root_command {
        return Err(anyhow!("expected root command `{root_command}`"));
    }

    let rest = args.get(2..).unwrap_or(&[]).to_vec();
    if rest.is_empty() {
        return Err(anyhow!(
            "expected `{root_command} list` or `{root_command} --target <target> --timeout-ms <u32>`"
        ));
    }

    if rest[0] == "list" {
        return observer_host_dispatch(provider, tests, &[String::from("list")]);
    }

    if rest[0] == "run" || rest[0] == "observe" {
        return observer_host_dispatch(provider, tests, &rest);
    }

    let mut routed = vec![String::from("observe")];
    routed.extend(rest);
    observer_host_dispatch(provider, tests, &routed)
}

fn validate_provider(provider: &str) -> Result<()> {
    if provider.trim().is_empty() {
        return Err(anyhow!("provider id must not be empty"));
    }
    Ok(())
}

fn write_json<T>(value: &T) -> Result<()>
where
    T: Serialize,
{
    println!("{}", serde_json::to_string(value)?);
    Ok(())
}

fn telemetry_payload_from_outcome(outcome: TestOutcome) -> Option<Vec<TelemetryPayloadEntry>> {
    let mut telemetry = Vec::new();
    for entry in outcome.telemetry {
        telemetry.push(convert_telemetry(entry));
    }
    if telemetry.is_empty() {
        None
    } else {
        Some(telemetry)
    }
}

fn convert_telemetry(entry: RustTelemetryEntry) -> TelemetryPayloadEntry {
    let value = match entry.value {
        RustTelemetryValue::Metric(value) => TelemetryPayloadValue::Metric { value },
        RustTelemetryValue::Vector(values) => TelemetryPayloadValue::Vector { values },
        RustTelemetryValue::Tag(value) => TelemetryPayloadValue::Tag { value },
    };

    TelemetryPayloadEntry {
        name: entry.name,
        unit: None,
        value,
    }
}

fn panic_payload_message(payload: &Box<dyn std::any::Any + Send>) -> String {
    if let Some(failure) = payload.downcast_ref::<ExpectationFailure>() {
        return failure.message().to_owned();
    }
    if let Some(message) = payload.downcast_ref::<&str>() {
        return format!("panic: {message}");
    }
    if let Some(message) = payload.downcast_ref::<String>() {
        return format!("panic: {message}");
    }
    String::from("panic: non-string payload")
}

#[macro_export]
macro_rules! describe {
    ($label:expr, $body:block) => {{
        $crate::describe($label, || $body)
    }};
}

#[macro_export]
macro_rules! test {
    ($title:expr, $function:expr $(,)?) => {{
        $crate::register_test($title, None, $function)
    }};
    ($title:expr, id = $id:expr, $function:expr $(,)?) => {{
        $crate::register_test($title, Some($id), $function)
    }};
}

#[macro_export]
macro_rules! it {
    ($title:expr, $function:expr $(,)?) => {{
        $crate::register_test($title, None, $function)
    }};
    ($title:expr, id = $id:expr, $function:expr $(,)?) => {{
        $crate::register_test($title, Some($id), $function)
    }};
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn collect_tests_derives_and_sorts_names() {
        let tests = collect_tests(|| {
            describe("pkg", || {
                test!("beta test", |ctx| {
                    ctx.stdout("beta\n");
                });
                test!("alpha test", id = "pkg/alpha", |ctx| {
                    ctx.stdout("alpha\n");
                });
            });
        })
        .expect("collection should validate");

        assert_eq!(tests.len(), 2);
        assert_eq!(tests[0].canonical_name(), "pkg :: beta test");
        assert_eq!(tests[1].canonical_name(), "pkg/alpha");
    }

    #[test]
    fn duplicate_derived_titles_gain_occurrence_suffixes() {
        let tests = collect_tests(|| {
            describe("pkg", || {
                test!("same", |_ctx| {});
                test!("same", |_ctx| {});
            });
        })
        .expect("collection should validate");

        assert_eq!(tests[0].canonical_name(), "pkg :: same");
        assert_eq!(tests[1].canonical_name(), "pkg :: same #2");
    }

    #[test]
    fn explicit_duplicate_ids_are_rejected() {
        let error = collect_tests(|| {
            describe("pkg", || {
                test!("one", id = "pkg/dup", |_ctx| {});
                test!("two", id = "pkg/dup", |_ctx| {});
            });
        })
        .expect_err("duplicate ids should fail");

        assert_eq!(error, RegistryError::DuplicateCanonicalName(String::from("pkg/dup")));
    }

    #[test]
    fn run_test_converts_expectation_failure_into_stderr() {
        let test = define_test(
            "smoke",
            |_ctx| {
                expect(false).to_be_truthy();
            },
            TestOptions::default(),
        )
        .expect("define_test should validate");

        let outcome = run_test(&test);
        assert_eq!(outcome.exit, 1);
        assert_eq!(String::from_utf8(outcome.err).expect("stderr should be utf8"), "expected false to be truthy\n");
    }

    #[test]
    fn list_payload_is_sorted() {
        let tests = vec![
            define_test("b", |_ctx| {}, TestOptions { id: Some("pkg/b") }).expect("valid"),
            define_test("a", |_ctx| {}, TestOptions { id: Some("pkg/a") }).expect("valid"),
        ];

        let payload = list_payload("rust", &tests).expect("list payload should succeed");
        assert_eq!(payload.tests[0].name, "pkg/a");
        assert_eq!(payload.tests[1].name, "pkg/b");
    }

    #[test]
    fn run_payload_encodes_output_and_telemetry() {
        let test = define_test(
            "smoke",
            |ctx| {
                ctx.stdout("ok\n");
                assert!(ctx.observe().metric("wall_time_ns", 42.0));
            },
            TestOptions { id: Some("pkg/smoke") },
        )
        .expect("valid test");

        let payload = run_payload("rust", &[test], "pkg/smoke", 1000).expect("run payload should succeed");
        assert_eq!(payload.exit, 0);
        assert_eq!(payload.out_b64, "b2sK");
        assert!(payload.telemetry.is_some());
    }
}