Files
dagx/tests/tracing/with_tracing.rs
T

81 lines
1.7 KiB
Rust

//! Tests with tracing feature enabled
use dagx::{task, DagRunner};
use tracing_subscriber::{fmt, EnvFilter};
struct Value(i32);
#[task]
impl Value {
async fn run(&self) -> i32 {
self.0
}
}
struct Add;
#[task]
impl Add {
async fn run(a: &i32, b: &i32) -> i32 {
a + b
}
}
struct Multiply;
#[task]
impl Multiply {
async fn run(a: &i32, b: &i32) -> i32 {
a * b
}
}
#[tokio::test(flavor = "multi_thread")]
async fn test_tracing_with_subscriber() {
// Initialize tracing subscriber for this test
let _ = fmt()
.with_env_filter(EnvFilter::from_default_env())
.with_test_writer()
.try_init();
let mut dag = DagRunner::new();
let a = dag.add_task(Value(2));
let b = dag.add_task(Value(3));
let sum = dag.add_task(Add).depends_on((&a, &b));
let mut output = dag
.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
.await
.unwrap();
assert_eq!(output.get(sum), 5);
}
#[tokio::test(flavor = "multi_thread")]
async fn test_tracing_inline_execution() {
// Initialize tracing subscriber
let _ = fmt()
.with_env_filter(EnvFilter::from_default_env())
.with_test_writer()
.try_init();
let mut dag = DagRunner::new();
// Create a linear chain (should trigger inline execution)
let a = dag.add_task(Value(1));
let b = dag.add_task(Add).depends_on((&a, &a));
let c = dag.add_task(Multiply).depends_on((&b, &b));
let mut output = dag
.run(|fut| async move { tokio::spawn(fut).await.unwrap() })
.await
.unwrap();
// a = 1
// b = 1 + 1 = 2
// c = 2 * 2 = 4
assert_eq!(output.get(c), 4);
}