rust/bt-common/src/debug_command.rs

// Copyright 2023 Google LLC
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Debug command traits and helpers for defining commands for integration
//! into a debug tool.
use futures::channel::mpsc::{unbounded, UnboundedReceiver, UnboundedSender};
use futures::task::AtomicWaker;
use futures::{FutureExt, Stream, StreamExt};
use log::LevelFilter;
use std::collections::HashMap;
use std::pin::Pin;
use std::str::FromStr;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};

#[must_use = "if unused the previous log level will immediately be restored"]
pub struct ScopedVerbosityGuard {
    prev_level: LevelFilter,
}

impl ScopedVerbosityGuard {
    pub fn new(scoped_level: LevelFilter) -> Self {
        let prev_level = log::max_level();
        log::set_max_level(scoped_level);
        Self { prev_level }
    }
}

impl Drop for ScopedVerbosityGuard {
    fn drop(&mut self) {
        log::set_max_level(self.prev_level);
    }
}

pub fn parse_level_filter(s: &str) -> Option<LevelFilter> {
    s.trim().parse().ok()
}

pub fn extract_verbosity_flags(args: Vec<String>) -> (Option<LevelFilter>, Vec<String>) {
    let mut level = None;
    let mut clean_args = Vec::new();
    let mut iter = args.into_iter().peekable();

    while let Some(arg) = iter.next() {
        match arg.as_str() {
            "-v" | "--verbose" => {
                level = Some(LevelFilter::Debug);
            }
            "-vv" | "--trace" => {
                level = Some(LevelFilter::Trace);
            }
            "-q" | "--quiet" => {
                level = Some(LevelFilter::Warn);
            }
            "--log-level" => {
                if let Some(next) = iter.peek() {
                    if let Some(lvl) = parse_level_filter(next) {
                        level = Some(lvl);
                        iter.next();
                        continue;
                    }
                }
                clean_args.push(arg);
            }
            _ if arg.starts_with("--log-level=") => {
                let val = &arg["--log-level=".len()..];
                if let Some(lvl) = parse_level_filter(val) {
                    level = Some(lvl);
                } else {
                    clean_args.push(arg);
                }
            }
            _ => {
                clean_args.push(arg);
            }
        }
    }

    (level, clean_args)
}

pub fn handle_verbose_command(args: &[String]) {
    match args.len() {
        0 => {
            println!("Current log level: {}", log::max_level());
        }
        1 => {
            if let Some(level) = parse_level_filter(&args[0]) {
                log::set_max_level(level);
                println!("Log level set to {}", level);
            } else {
                eprintln!(
                    "Invalid log level '{}'. Valid levels: off, error, warn, info, debug, trace",
                    args[0]
                );
            }
        }
        _ => {
            println!("Usage: verbose [off|error|warn|info|debug|trace]");
        }
    }
}

/// A CommandSet is a set of commands (usually an enum) that each represent an
/// action that can be performed.  i.e. 'list', 'volume' etc.  Each command can
/// take zero or more arguments and have zero or more flags.
/// Typically an Enum of commands would implement CommandSet trait.
pub trait CommandSet: FromStr + ::core::fmt::Display {
    /// Returns a vector of strings that are the commands supported by this.
    fn variants() -> Vec<String>;

    /// Returns a string listing the arguments that this command takes, in <>
    /// brackets
    fn arguments(&self) -> &'static str;

    /// Returns a string displaying the flags that this command supports, in []
    /// brackets
    fn flags(&self) -> &'static str;

    /// Returns a short description of this command
    fn desc(&self) -> &'static str;

    /// Help string for this variant (build from Display, arguments and flags by
    /// default)
    fn help_simple(&self) -> String {
        format!("{self} {} {} -- {}", self.flags(), self.arguments(), self.desc())
    }

    /// Possibly multi-line help string for all variants of this set.
    fn help_all() -> String {
        let mut help: String = Self::variants()
            .into_iter()
            .filter_map(|s| FromStr::from_str(&s).ok())
            .map(|s: Self| format!("{}\n", s))
            .collect();
        if !help.is_empty() {
            help.push('\n');
        }
        help.push_str("Global verbosity flags: [-v|--verbose] [-vv|--trace] [-q|--quiet] [--log-level=<level>]\n");
        help
    }
}

/// A top-level CLI command: either a framework meta-command (e.g. `verbose`)
/// or a profile-specific domain command.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CliCommand<T> {
    Verbose,
    Domain(T),
}

impl<T> From<T> for CliCommand<T> {
    fn from(cmd: T) -> Self {
        CliCommand::Domain(cmd)
    }
}

impl<T: CommandSet> CommandSet for CliCommand<T> {
    fn variants() -> Vec<String> {
        let mut variants = T::variants();
        variants.push("verbose".to_string());
        variants
    }

    fn arguments(&self) -> &'static str {
        match self {
            CliCommand::Domain(cmd) => cmd.arguments(),
            CliCommand::Verbose => "[level] ",
        }
    }

    fn flags(&self) -> &'static str {
        match self {
            CliCommand::Domain(cmd) => cmd.flags(),
            CliCommand::Verbose => "",
        }
    }

    fn desc(&self) -> &'static str {
        match self {
            CliCommand::Domain(cmd) => cmd.desc(),
            CliCommand::Verbose => "Get or set the logging verbosity level",
        }
    }
}

impl<T: ::core::fmt::Display> ::core::fmt::Display for CliCommand<T> {
    fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
        match self {
            CliCommand::Domain(cmd) => write!(f, "{cmd}"),
            CliCommand::Verbose => write!(f, "verbose"),
        }
    }
}

impl<T: FromStr> FromStr for CliCommand<T> {
    type Err = T::Err;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s == "verbose" {
            Ok(CliCommand::Verbose)
        } else {
            T::from_str(s).map(CliCommand::Domain)
        }
    }
}

/// Macro to help build CommandSets
#[macro_export]
macro_rules! gen_commandset {
    ($name:ident {
        $($variant:ident = ($val:expr, [$($flag:expr),*], [$($arg:expr),*], $help:expr)),*,
    }) => {
        /// Enum of all possible commands
        #[derive(PartialEq, Debug, Clone, Copy)]
        pub enum $name {
            $($variant,)*
        }

        impl CommandSet for $name {
            fn variants() -> Vec<String> {
                let mut variants = Vec::new();
                $(variants.push($val.to_string());)*
                variants
            }

            fn arguments(&self) -> &'static str {
                match self {
                    $(
                        $name::$variant => concat!($("<", $arg, "> ",)*),
                    )*
                }
            }

            fn flags(&self) -> &'static str {
                match self {
                    $(
                        $name::$variant => concat!($("[", $flag, "] ",)*),
                    )*
                }
            }

            fn desc(&self) -> &'static str {
                match self {
                    $(
                        $name::$variant => $help,
                    )*
                }
            }
        }

        impl ::core::fmt::Display for $name {
            fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
                match *self {
                    $($name::$variant => write!(f, $val),)*
                }
            }
        }

        impl ::std::str::FromStr for $name {
            type Err = ();

            fn from_str(s: &str) -> Result<$name, ()> {
                match s {
                    $($val => Ok($name::$variant),)*
                    _ => Err(()),
                }
            }
        }
    }
}

/// An event emitted by a [`CommandRunner`] through its
/// [`CommandRunner::event_stream`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RunnerEvent<E = ()> {
    /// An arbitrary human-readable message from the runner or its background
    /// tasks to be displayed by the harness.
    Message(String),
    /// A structured domain event specific to this runner that the harness or a
    /// parent composite tool can react to programmatically.
    Event(E),
}

impl<E> RunnerEvent<E> {
    pub fn message(&self) -> Option<&str> {
        match self {
            Self::Message(s) => Some(s),
            Self::Event(_) => None,
        }
    }

    pub fn event(&self) -> Option<&E> {
        match self {
            Self::Message(_) => None,
            Self::Event(e) => Some(e),
        }
    }

    pub fn into_event(self) -> Option<E> {
        match self {
            Self::Message(_) => None,
            Self::Event(e) => Some(e),
        }
    }

    pub fn into_message(self) -> Option<String> {
        match self {
            Self::Message(s) => Some(s),
            Self::Event(_) => None,
        }
    }

    pub fn map_event<F, T>(self, f: F) -> RunnerEvent<T>
    where
        F: FnOnce(E) -> T,
    {
        match self {
            Self::Message(m) => RunnerEvent::Message(m),
            Self::Event(e) => RunnerEvent::Event(f(e)),
        }
    }
}

impl<E> From<String> for RunnerEvent<E> {
    fn from(s: String) -> Self {
        Self::Message(s)
    }
}

impl<E> From<&str> for RunnerEvent<E> {
    fn from(s: &str) -> Self {
        Self::Message(s.to_string())
    }
}

/// CommandRunner is used to perform a specific task based on the command set.
pub trait CommandRunner {
    type Set: CommandSet;
    type Error: ::std::error::Error;
    type Event;

    fn run_command(
        &self,
        cmd: Self::Set,
        args: Vec<String>,
    ) -> impl futures::Future<Output = Result<(), Self::Error>>;

    /// Returns a stream that drives background processing for this runner and
    /// yields display messages or structured events for the harness.
    ///
    /// Crates that do not need background tasks or event delivery can set
    /// `type Event = ();` and rely on this default implementation.
    fn event_stream(&self) -> impl Stream<Item = RunnerEvent<Self::Event>> + '_ {
        futures::stream::pending::<RunnerEvent<Self::Event>>()
    }

    fn run(
        &self,
        cmd: impl Into<CliCommand<Self::Set>>,
        args: Vec<String>,
    ) -> impl futures::Future<Output = Result<(), Self::Error>> {
        let cmd = cmd.into();
        async move {
            match cmd {
                CliCommand::Verbose => {
                    handle_verbose_command(&args);
                    Ok(())
                }
                CliCommand::Domain(cmd) => {
                    let (scoped_level, clean_args) = extract_verbosity_flags(args);
                    let _guard = scoped_level.map(ScopedVerbosityGuard::new);
                    self.run_command(cmd, clean_args).await
                }
            }
        }
    }
}

impl<R: CommandRunner> CommandRunner for &R {
    type Set = R::Set;
    type Error = R::Error;
    type Event = R::Event;

    fn run_command(
        &self,
        cmd: Self::Set,
        args: Vec<String>,
    ) -> impl futures::Future<Output = Result<(), Self::Error>> {
        (**self).run_command(cmd, args)
    }

    fn event_stream(&self) -> impl Stream<Item = RunnerEvent<Self::Event>> + '_ {
        (**self).event_stream()
    }
}

impl<R: CommandRunner> CommandRunner for Arc<R> {
    type Set = R::Set;
    type Error = R::Error;
    type Event = R::Event;

    fn run_command(
        &self,
        cmd: Self::Set,
        args: Vec<String>,
    ) -> impl futures::Future<Output = Result<(), Self::Error>> {
        (**self).run_command(cmd, args)
    }

    fn event_stream(&self) -> impl Stream<Item = RunnerEvent<Self::Event>> + '_ {
        (**self).event_stream()
    }
}

type BoxedStream<E> = Pin<Box<dyn Stream<Item = RunnerEvent<E>> + Send + 'static>>;

struct StreamsState<E> {
    keyed: HashMap<String, BoxedStream<E>>,
    anonymous: Vec<BoxedStream<E>>,
}

impl<E> Default for StreamsState<E> {
    fn default() -> Self {
        Self { keyed: HashMap::new(), anonymous: Vec::new() }
    }
}

impl<E: Send + 'static> Stream for StreamsState<E> {
    type Item = RunnerEvent<E>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut ready = None;
        self.keyed.retain(|_key, stream| {
            if ready.is_some() {
                return true;
            }
            match stream.as_mut().poll_next(cx) {
                Poll::Ready(Some(item)) => {
                    ready = Some(item);
                    true
                }
                Poll::Ready(None) => false,
                Poll::Pending => true,
            }
        });
        if let Some(item) = ready {
            return Poll::Ready(Some(item));
        }

        self.anonymous.retain_mut(|stream| {
            if ready.is_some() {
                return true;
            }
            match stream.as_mut().poll_next(cx) {
                Poll::Ready(Some(item)) => {
                    ready = Some(item);
                    true
                }
                Poll::Ready(None) => false,
                Poll::Pending => true,
            }
        });
        if let Some(item) = ready {
            return Poll::Ready(Some(item));
        }

        Poll::Pending
    }
}

struct RunnerBackgroundInner<E: Send + 'static> {
    msg_sender: UnboundedSender<RunnerEvent<E>>,
    msg_receiver: Mutex<UnboundedReceiver<RunnerEvent<E>>>,
    streams: Mutex<StreamsState<E>>,
    waker: AtomicWaker,
}

/// Helper for managing background streams, display messages, and structured
/// events for a [`CommandRunner`].
#[derive(Clone)]
pub struct RunnerBackground<E: Send + 'static = ()> {
    inner: Arc<RunnerBackgroundInner<E>>,
}

impl<E: Send + 'static> RunnerBackground<E> {
    pub fn new() -> Self {
        let (msg_sender, msg_receiver) = unbounded();
        Self {
            inner: Arc::new(RunnerBackgroundInner {
                msg_sender,
                msg_receiver: Mutex::new(msg_receiver),
                streams: Mutex::new(StreamsState::default()),
                waker: AtomicWaker::new(),
            }),
        }
    }

    /// Queues a displayable text message to be emitted by
    /// [`Self::event_stream`].
    pub fn println(&self, msg: impl Into<String>) {
        let _ = self.inner.msg_sender.unbounded_send(RunnerEvent::Message(msg.into()));
        self.inner.waker.wake();
    }

    /// Queues a structured domain event to be emitted by
    /// [`Self::event_stream`].
    pub fn send(&self, event: E) {
        let _ = self.inner.msg_sender.unbounded_send(RunnerEvent::Event(event));
        self.inner.waker.wake();
    }

    /// Returns an [`UnboundedSender`] that converts sent items into
    /// [`RunnerEvent::Event`].
    pub fn event_sender(&self) -> UnboundedSender<E> {
        let (tx, rx) = unbounded::<E>();
        self.spawn_stream(rx.map(RunnerEvent::Event));
        tx
    }

    /// Dynamically registers or replaces a background stream under a unique
    /// key. If a stream was already registered under `key`, the previous
    /// stream is cancelled.
    pub fn set_stream<K: Into<String>>(
        &self,
        key: K,
        stream: impl Stream<Item = RunnerEvent<E>> + Send + 'static,
    ) {
        let mut streams = self.inner.streams.lock().unwrap();
        streams.keyed.insert(key.into(), Box::pin(stream));
        drop(streams);
        self.inner.waker.wake();
    }

    /// Cancels the stream registered under `key`, if present. Returns `true` if
    /// a stream was cancelled.
    pub fn cancel_stream(&self, key: &str) -> bool {
        let mut streams = self.inner.streams.lock().unwrap();
        let removed = streams.keyed.remove(key).is_some();
        drop(streams);
        if removed {
            self.inner.waker.wake();
        }
        removed
    }

    /// Spawns an anonymous background stream that runs until completion.
    pub fn spawn_stream(&self, stream: impl Stream<Item = RunnerEvent<E>> + Send + 'static) {
        let mut streams = self.inner.streams.lock().unwrap();
        streams.anonymous.push(Box::pin(stream));
        drop(streams);
        self.inner.waker.wake();
    }

    /// Returns `true` if a background stream is currently registered under
    /// `key`.
    pub fn has_stream(&self, key: &str) -> bool {
        self.inner.streams.lock().unwrap().keyed.contains_key(key)
    }

    /// Returns the total number of active background streams (keyed and
    /// anonymous).
    pub fn active_stream_count(&self) -> usize {
        let streams = self.inner.streams.lock().unwrap();
        streams.keyed.len() + streams.anonymous.len()
    }

    /// Returns a stream yielding messages and events produced by this runner.
    pub fn event_stream(&self) -> RunnerBackgroundStream<E> {
        RunnerBackgroundStream { inner: self.inner.clone() }
    }
}

impl<E: Send + 'static> RunnerBackgroundInner<E> {
    fn poll_next_event(&self, cx: &mut Context<'_>) -> Poll<Option<RunnerEvent<E>>> {
        self.waker.register(cx.waker());

        // 1. Drain directly queued messages/events from msg_receiver first.
        if let Poll::Ready(Some(item)) = self.msg_receiver.lock().unwrap().poll_next_unpin(cx) {
            return Poll::Ready(Some(item));
        }

        // 2. Poll registered background streams.
        if let Poll::Ready(Some(item)) = self.streams.lock().unwrap().poll_next_unpin(cx) {
            return Poll::Ready(Some(item));
        }

        Poll::Pending
    }
}

impl<E: Send + 'static> Default for RunnerBackground<E> {
    fn default() -> Self {
        Self::new()
    }
}

/// Stream returned by [`RunnerBackground::event_stream`].
pub struct RunnerBackgroundStream<E: Send + 'static = ()> {
    inner: Arc<RunnerBackgroundInner<E>>,
}

impl<E: Send + 'static> Stream for RunnerBackgroundStream<E> {
    type Item = RunnerEvent<E>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.poll_next_event(cx)
    }
}

/// A harness for driving a [`CommandRunner`] alongside its background event
/// stream.
pub struct DebugHarness<R: CommandRunner> {
    runner: R,
}

impl<R: CommandRunner> DebugHarness<R> {
    pub fn new(runner: R) -> Self {
        Self { runner }
    }

    pub fn runner(&self) -> &R {
        &self.runner
    }

    pub fn runner_mut(&mut self) -> &mut R {
        &mut self.runner
    }

    pub fn into_runner(self) -> R {
        self.runner
    }

    pub fn event_stream(&self) -> impl Stream<Item = RunnerEvent<R::Event>> + '_ {
        self.runner.event_stream()
    }

    pub async fn run(
        &self,
        cmd: impl Into<CliCommand<R::Set>>,
        args: Vec<String>,
    ) -> Result<(), R::Error> {
        self.runner.run(cmd, args).await
    }

    async fn drive_events<Fut, F>(&self, future: Fut, mut on_event: F) -> Fut::Output
    where
        Fut: futures::Future,
        F: FnMut(RunnerEvent<R::Event>),
    {
        let stream = self.runner.event_stream().fuse();
        futures::pin_mut!(stream);
        let future = future.fuse();
        futures::pin_mut!(future);

        loop {
            futures::select! {
                event = stream.select_next_some() => on_event(event),
                output = future => {
                    while let Some(Some(event)) = stream.next().now_or_never() {
                        on_event(event);
                    }
                    return output;
                }
            }
        }
    }

    /// Concurrently drives the runner's background `event_stream` while
    /// executing `future`. Messages are passed to `on_message` and events
    /// are passed to `on_event`.
    pub async fn drive_with<Fut, FM, FE>(
        &self,
        future: Fut,
        mut on_message: FM,
        mut on_event: FE,
    ) -> Fut::Output
    where
        Fut: futures::Future,
        FM: FnMut(String),
        FE: FnMut(R::Event),
    {
        self.drive_events(future, |event| match event {
            RunnerEvent::Message(msg) => on_message(msg),
            RunnerEvent::Event(evt) => on_event(evt),
        })
        .await
    }

    /// Concurrently executes `future` while collecting all `RunnerEvent`s
    /// emitted during its execution.
    pub async fn run_collecting_events<Fut>(
        &self,
        future: Fut,
    ) -> (Fut::Output, Vec<RunnerEvent<R::Event>>)
    where
        Fut: futures::Future,
    {
        let mut events = Vec::new();
        let output = self.drive_events(future, |event| events.push(event)).await;
        (output, events)
    }

    /// Concurrently executes a command while collecting all `RunnerEvent`s
    /// emitted during its execution.
    pub async fn run_command_collecting(
        &self,
        cmd: impl Into<CliCommand<R::Set>>,
        args: Vec<String>,
    ) -> (Result<(), R::Error>, Vec<RunnerEvent<R::Event>>) {
        self.run_collecting_events(self.run(cmd, args)).await
    }
}

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

    #[test]
    fn gen_commandset_simple() {
        gen_commandset! {
            TestCmd {
                One = ("one", [], [], "First Command"),
                WithFlags = ("with-flags", ["-1","-2"], [], "Command with flags"),
                WithArgs = ("with-args", [], ["arg", "two"], "Command with args"),
                WithBoth = ("with-both", ["-w"], ["simple"], "Command with both flags and args"),
            }
        }

        let cmd: TestCmd = "one".parse().unwrap();
        assert_eq!(cmd, TestCmd::One);

        let cmd2: TestCmd = "with-flags".parse().unwrap();
        assert_eq!(cmd2.arguments(), "");
        assert_eq!(cmd2.flags(), "[-1] [-2] ");

        let verbose_cmd: CliCommand<TestCmd> = "verbose".parse().unwrap();
        assert_eq!(verbose_cmd, CliCommand::Verbose);
        assert_eq!(verbose_cmd.flags(), "");
        assert_eq!(verbose_cmd.to_string(), "verbose");

        let domain_cmd: CliCommand<TestCmd> = "one".parse().unwrap();
        assert_eq!(domain_cmd, CliCommand::Domain(TestCmd::One));

        assert!(!TestCmd::variants().contains(&"verbose".to_string()));
        assert!(CliCommand::<TestCmd>::variants().contains(&"verbose".to_string()));
        assert!(TestCmd::help_all().contains("Global verbosity flags:"));
        assert!(CliCommand::<TestCmd>::help_all().contains("Global verbosity flags:"));
    }

    #[test]
    fn test_extract_verbosity_flags() {
        // -v / --verbose
        let (lvl, args) = extract_verbosity_flags(vec!["-v".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Debug));
        assert!(args.is_empty());

        let (lvl, args) = extract_verbosity_flags(vec!["--verbose".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Debug));
        assert!(args.is_empty());

        // -vv / --trace
        let (lvl, args) = extract_verbosity_flags(vec!["-vv".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Trace));
        assert!(args.is_empty());

        let (lvl, args) = extract_verbosity_flags(vec!["--trace".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Trace));
        assert!(args.is_empty());

        // -q / --quiet
        let (lvl, args) = extract_verbosity_flags(vec!["-q".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Warn));
        assert!(args.is_empty());

        let (lvl, args) = extract_verbosity_flags(vec!["--quiet".to_string()]);
        assert_eq!(lvl, Some(LevelFilter::Warn));
        assert!(args.is_empty());

        // --log-level=<level> and --log-level <level> (case-insensitive)
        for (name, expected) in [
            ("off", LevelFilter::Off),
            ("error", LevelFilter::Error),
            ("warn", LevelFilter::Warn),
            ("info", LevelFilter::Info),
            ("debug", LevelFilter::Debug),
            ("trace", LevelFilter::Trace),
            ("DEBUG", LevelFilter::Debug),
            ("Info", LevelFilter::Info),
        ] {
            let (lvl, args) = extract_verbosity_flags(vec![format!("--log-level={name}")]);
            assert_eq!(lvl, Some(expected));
            assert!(args.is_empty());

            let (lvl, args) =
                extract_verbosity_flags(vec!["--log-level".to_string(), name.to_string()]);
            assert_eq!(lvl, Some(expected));
            assert!(args.is_empty());
        }

        // Mixed arguments
        let input = vec![
            "action".to_string(),
            "-v".to_string(),
            "--custom-flag".to_string(),
            "target".to_string(),
            "--log-level=warn".to_string(),
            "extra".to_string(),
        ];
        let (lvl, args) = extract_verbosity_flags(input);
        assert_eq!(lvl, Some(LevelFilter::Warn));
        assert_eq!(args, vec!["action", "--custom-flag", "target", "extra"]);

        // Invalid log level flag should not be extracted
        let (lvl, args) =
            extract_verbosity_flags(vec!["--log-level=invalid".to_string(), "arg".to_string()]);
        assert_eq!(lvl, None);
        assert_eq!(args, vec!["--log-level=invalid", "arg"]);
    }

    // Mutex to avoid parallel tests fighting over log levels
    static LOG_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(());

    #[test]
    fn test_scoped_verbosity_guard() {
        let _lock = LOG_MUTEX.lock().unwrap();
        log::set_max_level(LevelFilter::Info);

        futures::executor::block_on(async {
            assert_eq!(log::max_level(), LevelFilter::Info);
            {
                let _guard = ScopedVerbosityGuard::new(LevelFilter::Trace);
                assert_eq!(log::max_level(), LevelFilter::Trace);
                // Yield across async boundary
                futures::future::ready(()).await;
                assert_eq!(log::max_level(), LevelFilter::Trace);
            }
            assert_eq!(log::max_level(), LevelFilter::Info);
        });
    }

    #[test]
    fn test_persistent_verbosity_and_handle_command() {
        let _lock = LOG_MUTEX.lock().unwrap();
        log::set_max_level(LevelFilter::Debug);

        // handle_verbose_command with 0 args
        handle_verbose_command(&[]);
        assert_eq!(log::max_level(), LevelFilter::Debug);

        // handle_verbose_command with 1 valid arg
        handle_verbose_command(&["trace".to_string()]);
        assert_eq!(log::max_level(), LevelFilter::Trace);

        handle_verbose_command(&["WARN".to_string()]);
        assert_eq!(log::max_level(), LevelFilter::Warn);

        // handle_verbose_command with invalid arg (should not change)
        handle_verbose_command(&["invalid_level".to_string()]);
        assert_eq!(log::max_level(), LevelFilter::Warn);

        // handle_verbose_command with > 1 args (should not change)
        handle_verbose_command(&["info".to_string(), "extra".to_string()]);
        assert_eq!(log::max_level(), LevelFilter::Warn);

        // Reset back to Info
        log::set_max_level(LevelFilter::Info);
    }

    #[test]
    fn test_command_runner_run() {
        gen_commandset! {
            RunnerCmd {
                Action = ("action", [], ["target"], "Perform action"),
            }
        }

        struct TestRunner {
            executed: std::sync::Mutex<Vec<(RunnerCmd, Vec<String>)>>,
        }

        impl CommandRunner for TestRunner {
            type Set = RunnerCmd;
            type Error = std::io::Error;
            type Event = ();

            fn run_command(
                &self,
                cmd: Self::Set,
                args: Vec<String>,
            ) -> impl futures::Future<Output = Result<(), Self::Error>> {
                self.executed.lock().unwrap().push((cmd, args));
                futures::future::ready(Ok(()))
            }
        }

        let runner = TestRunner { executed: std::sync::Mutex::new(Vec::new()) };
        let _lock = LOG_MUTEX.lock().unwrap();

        // Running Verbose command should update persistent verbosity and NOT
        // invoke run_command
        log::set_max_level(LevelFilter::Info);
        futures::executor::block_on(runner.run(CliCommand::Verbose, vec!["trace".to_string()]))
            .unwrap();
        assert_eq!(log::max_level(), LevelFilter::Trace);
        assert!(runner.executed.lock().unwrap().is_empty());

        // Running normal command with verbosity flag strips flag and runs
        // command
        let args = vec!["-v".to_string(), "my_target".to_string()];
        futures::executor::block_on(runner.run(RunnerCmd::Action, args)).unwrap();
        let executed = runner.executed.lock().unwrap().clone();
        assert_eq!(executed.len(), 1);
        assert_eq!(executed[0], (RunnerCmd::Action, vec!["my_target".to_string()]));

        // After run completes, the guard has restored the previous level
        assert_eq!(log::max_level(), LevelFilter::Trace);

        // Reset persistent verbosity back to Info
        log::set_max_level(LevelFilter::Info);
    }

    #[test]
    fn test_command_runner_default_event_stream() {
        gen_commandset! {
            DefaultCmd {
                Action = ("action", [], [], "Default action"),
            }
        }

        struct SimpleRunner;
        impl CommandRunner for SimpleRunner {
            type Set = DefaultCmd;
            type Error = std::io::Error;
            type Event = ();

            fn run_command(
                &self,
                _cmd: Self::Set,
                _args: Vec<String>,
            ) -> impl futures::Future<Output = Result<(), Self::Error>> {
                futures::future::ready(Ok(()))
            }
        }

        let runner = SimpleRunner;
        let stream = runner.event_stream();
        futures::pin_mut!(stream);
        let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref());
        assert!(stream.as_mut().poll_next(&mut cx).is_pending());
    }

    #[test]
    fn test_runner_event() {
        let msg: RunnerEvent<u32> = RunnerEvent::Message("hello".to_string());
        assert_eq!(msg.message(), Some("hello"));
        assert_eq!(msg.event(), None);
        assert_eq!(msg.clone().into_event(), None);
        assert_eq!(msg.clone().into_message(), Some("hello".to_string()));

        let evt: RunnerEvent<u32> = RunnerEvent::Event(42);
        assert_eq!(evt.message(), None);
        assert_eq!(evt.event(), Some(&42));
        assert_eq!(evt.clone().into_event(), Some(42));
        assert_eq!(evt.into_message(), None);

        let str_msg: RunnerEvent<()> = "test msg".into();
        assert_eq!(str_msg, RunnerEvent::Message("test msg".to_string()));

        let string_msg: RunnerEvent<()> = String::from("string msg").into();
        assert_eq!(string_msg, RunnerEvent::Message("string msg".to_string()));

        let mapped = RunnerEvent::Event(10).map_event(|x| x * 2);
        assert_eq!(mapped, RunnerEvent::Event(20));

        let mapped_msg: RunnerEvent<i32> =
            RunnerEvent::Message("hi".to_string()).map_event(|x: u32| x as i32);
        assert_eq!(mapped_msg, RunnerEvent::Message("hi".to_string()));
    }

    #[test]
    fn test_runner_background_println_and_send() {
        futures::executor::block_on(async {
            let bg: RunnerBackground<u32> = RunnerBackground::new();
            bg.println("log message 1");
            bg.send(100);
            bg.println("log message 2");

            let stream = bg.event_stream();
            futures::pin_mut!(stream);

            assert_eq!(
                stream.next().await,
                Some(RunnerEvent::Message("log message 1".to_string()))
            );
            assert_eq!(stream.next().await, Some(RunnerEvent::Event(100)));
            assert_eq!(
                stream.next().await,
                Some(RunnerEvent::Message("log message 2".to_string()))
            );

            // When no more events, polling should return Pending.
            let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref());
            assert!(stream.as_mut().poll_next(&mut cx).is_pending());
        });
    }

    #[test]
    fn test_runner_background_event_sender() {
        futures::executor::block_on(async {
            let bg: RunnerBackground<String> = RunnerBackground::new();
            let sender = bg.event_sender();

            sender.unbounded_send("hello".to_string()).unwrap();
            sender.unbounded_send("world".to_string()).unwrap();

            let stream = bg.event_stream();
            futures::pin_mut!(stream);

            assert_eq!(stream.next().await, Some(RunnerEvent::Event("hello".to_string())));
            assert_eq!(stream.next().await, Some(RunnerEvent::Event("world".to_string())));

            drop(sender);
            // Polling after dropping sender should clean up the stream and
            // return Pending.
            let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref());
            assert!(stream.as_mut().poll_next(&mut cx).is_pending());
            assert_eq!(bg.active_stream_count(), 0);
        });
    }

    #[test]
    fn test_runner_background_keyed_streams_and_cancellation() {
        futures::executor::block_on(async {
            let bg: RunnerBackground<u32> = RunnerBackground::new();
            assert_eq!(bg.has_stream("discovery"), false);

            let (tx1, rx1) = futures::channel::mpsc::unbounded::<RunnerEvent<u32>>();
            bg.set_stream("discovery", rx1);
            assert_eq!(bg.has_stream("discovery"), true);
            assert_eq!(bg.active_stream_count(), 1);

            tx1.unbounded_send(RunnerEvent::Event(1)).unwrap();
            let stream = bg.event_stream();
            futures::pin_mut!(stream);
            assert_eq!(stream.next().await, Some(RunnerEvent::Event(1)));

            // Replacing the stream cancels the old one.
            let (tx2, rx2) = futures::channel::mpsc::unbounded::<RunnerEvent<u32>>();
            bg.set_stream("discovery", rx2);
            assert_eq!(bg.has_stream("discovery"), true);
            // Old tx1 should have receiver dropped / closed.
            assert!(tx1.is_closed());

            tx2.unbounded_send(RunnerEvent::Event(2)).unwrap();
            assert_eq!(stream.next().await, Some(RunnerEvent::Event(2)));

            // Cancelling the stream.
            assert!(bg.cancel_stream("discovery"));
            assert_eq!(bg.has_stream("discovery"), false);
            assert!(tx2.is_closed());
            assert!(!bg.cancel_stream("discovery"));
            assert_eq!(bg.active_stream_count(), 0);
        });
    }

    #[test]
    fn test_runner_background_spawn_stream_and_cleanup() {
        futures::executor::block_on(async {
            let bg: RunnerBackground<u32> = RunnerBackground::new();
            let stream_data = futures::stream::iter(vec![
                RunnerEvent::Message("from stream".to_string()),
                RunnerEvent::Event(99),
            ]);
            bg.spawn_stream(stream_data);
            assert_eq!(bg.active_stream_count(), 1);

            let stream = bg.event_stream();
            futures::pin_mut!(stream);

            assert_eq!(stream.next().await, Some(RunnerEvent::Message("from stream".to_string())));
            assert_eq!(stream.next().await, Some(RunnerEvent::Event(99)));

            // Polling again will observe the stream returned None and clean it
            // up.
            let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref());
            assert!(stream.as_mut().poll_next(&mut cx).is_pending());
            assert_eq!(bg.active_stream_count(), 0);
        });
    }

    #[test]
    fn test_runner_background_waker_waking() {
        let bg: RunnerBackground<u32> = RunnerBackground::new();
        let bg_clone = bg.clone();
        let (polled_tx, polled_rx) = futures::channel::oneshot::channel();

        let handle = std::thread::spawn(move || {
            futures::executor::block_on(polled_rx).unwrap();
            bg_clone.println("woken up");
            bg_clone.send(42);
        });

        futures::executor::block_on(async {
            let stream = bg.event_stream();
            futures::pin_mut!(stream);
            let mut polled_tx = Some(polled_tx);
            let first = futures::future::poll_fn(|cx| {
                let res = stream.as_mut().poll_next(cx);
                if let Some(tx) = polled_tx.take() {
                    assert!(res.is_pending());
                    tx.send(()).unwrap();
                }
                res
            })
            .await;
            assert_eq!(first, Some(RunnerEvent::Message("woken up".to_string())));
            assert_eq!(stream.next().await, Some(RunnerEvent::Event(42)));
        });

        handle.join().unwrap();
    }

    #[test]
    fn test_debug_harness() {
        gen_commandset! {
            HarnessCmd {
                Greet = ("greet", [], ["name"], "Greet a user"),
                Trigger = ("trigger", [], [], "Trigger background event"),
            }
        }

        struct HarnessRunner {
            bg: RunnerBackground<String>,
        }

        impl CommandRunner for HarnessRunner {
            type Set = HarnessCmd;
            type Error = std::io::Error;
            type Event = String;

            fn run_command(
                &self,
                cmd: Self::Set,
                args: Vec<String>,
            ) -> impl futures::Future<Output = Result<(), Self::Error>> {
                match cmd {
                    HarnessCmd::Greet => {
                        let name = args.first().cloned().unwrap_or_else(|| "world".to_string());
                        self.bg.println(format!("Hello, {}!", name));
                    }
                    HarnessCmd::Trigger => {
                        self.bg.send("event_fired".to_string());
                    }
                }
                futures::future::ready(Ok(()))
            }

            fn event_stream(&self) -> impl Stream<Item = RunnerEvent<Self::Event>> + '_ {
                self.bg.event_stream()
            }
        }

        let runner = HarnessRunner { bg: RunnerBackground::new() };
        let harness = DebugHarness::new(runner);

        futures::executor::block_on(async {
            // Test run_command_collecting
            let (res, events) =
                harness.run_command_collecting(HarnessCmd::Greet, vec!["Alice".to_string()]).await;
            assert!(res.is_ok());
            assert_eq!(events, vec![RunnerEvent::Message("Hello, Alice!".to_string())]);

            // Test drive_with with Trigger
            let mut messages = Vec::new();
            let mut domain_events = Vec::new();
            harness
                .drive_with(
                    harness.run(HarnessCmd::Trigger, vec![]),
                    |msg| messages.push(msg),
                    |evt| domain_events.push(evt),
                )
                .await
                .unwrap();

            assert!(messages.is_empty());
            assert_eq!(domain_events, vec!["event_fired".to_string()]);
        });
    }

    #[test]
    fn test_debug_harness_terminating_stream() {
        gen_commandset! {
            FiniteCmd {
                Action = ("action", [], [], "Finite action"),
            }
        }

        struct FiniteStreamRunner {
            stream_done_tx: std::sync::Mutex<Option<futures::channel::oneshot::Sender<()>>>,
        }
        impl CommandRunner for FiniteStreamRunner {
            type Set = FiniteCmd;
            type Error = std::io::Error;
            type Event = u32;

            fn run_command(
                &self,
                _cmd: Self::Set,
                _args: Vec<String>,
            ) -> impl futures::Future<Output = Result<(), Self::Error>> {
                futures::future::ready(Ok(()))
            }

            fn event_stream(&self) -> impl Stream<Item = RunnerEvent<Self::Event>> + '_ {
                let mut done_tx = self.stream_done_tx.lock().unwrap().take();
                futures::stream::iter(vec![
                    RunnerEvent::Message("finite msg".to_string()),
                    RunnerEvent::Event(42),
                ])
                .chain(futures::stream::poll_fn(move |_| {
                    if let Some(tx) = done_tx.take() {
                        let _ = tx.send(());
                    }
                    Poll::Ready(None)
                }))
            }
        }

        let (stream_done_tx, stream_done_rx) = futures::channel::oneshot::channel::<()>();
        let harness = DebugHarness::new(FiniteStreamRunner {
            stream_done_tx: std::sync::Mutex::new(Some(stream_done_tx)),
        });
        futures::executor::block_on(async {
            let (tx, rx) = futures::channel::oneshot::channel::<()>();
            let handle = std::thread::spawn(move || {
                futures::executor::block_on(stream_done_rx).unwrap();
                let _ = tx.send(());
            });

            // The stream terminates after 2 items while rx is still pending.
            // run_collecting_events should not busy-loop and should cleanly
            // complete when rx resolves.
            let (output, events) = harness.run_collecting_events(rx).await;
            assert!(output.is_ok());
            assert_eq!(
                events,
                vec![RunnerEvent::Message("finite msg".to_string()), RunnerEvent::Event(42),]
            );
            handle.join().unwrap();
        });
    }

    #[test]
    fn test_runner_background_stream_static() {
        let bg: RunnerBackground<u32> = RunnerBackground::new();
        bg.send(77);
        let stream = bg.event_stream();

        fn assert_static<T: 'static>(_: &T) {}
        assert_static(&stream);

        futures::executor::block_on(async {
            futures::pin_mut!(stream);
            assert_eq!(stream.next().await, Some(RunnerEvent::Event(77)));
        });
    }
}