From 4867098017156a341d5e62469d2cf579b370ffb5 Mon Sep 17 00:00:00 2001
From: jfldde <168934971+jfldde@users.noreply.github.com>
Date: Fri, 11 Oct 2024 10:13:11 +0100
Subject: [PATCH] Set up tracing subscriber on TestFramework init

---
 Cargo.toml          |  2 ++
 src/batch_prover.rs |  4 ++--
 src/framework.rs    | 35 ++++++++++++++++++++++++++++++++++-
 3 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 353b79c..eb40d0b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -20,6 +20,8 @@ tempfile = "3.8"
 tokio = { version = "1.39", features = ["full"] }
 toml = "0.8.0"
 tracing = { version = "0.1.40", default-features = false }
+tracing-subscriber = { version = "0.3.17", features = ["env-filter", "json", "fmt"] }
+log-panics = { version = "2", features = ["with-backtrace"] }
 
 # Citrea dependencies
 sov-ledger-rpc = { git = "https://github.com/chainwayxyz/citrea", rev = "82bf52d", default-features = false, features = ["client"] }
diff --git a/src/batch_prover.rs b/src/batch_prover.rs
index f7b427c..0a3ed11 100644
--- a/src/batch_prover.rs
+++ b/src/batch_prover.rs
@@ -2,7 +2,7 @@ use std::time::SystemTime;
 
 use anyhow::bail;
 use tokio::time::{sleep, Duration};
-use tracing::debug;
+use tracing::trace;
 
 use super::{config::FullBatchProverConfig, Result};
 use crate::node::Node;
@@ -14,7 +14,7 @@ impl BatchProver {
         let start = SystemTime::now();
         let timeout = timeout.unwrap_or(Duration::from_secs(600));
         loop {
-            debug!("Waiting for batch prover height {}", height);
+            trace!("Waiting for batch prover height {}", height);
             let latest_block = self.client.ledger_get_last_scanned_l1_height().await?;
 
             if latest_block >= height {
diff --git a/src/framework.rs b/src/framework.rs
index d667b6d..8fd74d9 100644
--- a/src/framework.rs
+++ b/src/framework.rs
@@ -1,4 +1,7 @@
-use std::{future::Future, sync::Arc};
+use std::{
+    future::Future,
+    sync::{Arc, Once},
+};
 
 use bitcoincore_rpc::RpcApi;
 use tracing::{debug, info};
@@ -53,6 +56,8 @@ async fn create_optional<T>(pred: bool, f: impl Future<Output = Result<T>>) -> R
 
 impl TestFramework {
     pub async fn new(config: TestConfig) -> Result<Self> {
+        setup_logging();
+
         anyhow::ensure!(
             config.test_case.n_nodes > 0,
             "At least one bitcoin node has to be running"
@@ -219,3 +224,31 @@ impl TestFramework {
         Ok(())
     }
 }
+
+use tracing_subscriber::layer::SubscriberExt;
+use tracing_subscriber::util::SubscriberInitExt;
+use tracing_subscriber::{fmt, EnvFilter};
+
+static INIT: Once = Once::new();
+
+fn setup_logging() {
+    INIT.call_once(|| {
+        let env_filter = EnvFilter::try_from_default_env()
+            .or_else(|_| EnvFilter::try_new("citrea_e2e=info"))
+            .unwrap();
+
+        if std::env::var("JSON_LOGS").is_ok() {
+            tracing_subscriber::registry()
+                .with(fmt::layer().json())
+                .with(env_filter)
+                .init();
+        } else {
+            tracing_subscriber::registry()
+                .with(fmt::layer())
+                .with(env_filter)
+                .init();
+        }
+
+        log_panics::init();
+    });
+}