From d3b6daf699c8e897aab2c86f8917c43c886ed32d Mon Sep 17 00:00:00 2001
From: jfldde <168934971+jfldde@users.noreply.github.com>
Date: Fri, 11 Oct 2024 16:22:26 +0100
Subject: [PATCH] Set up tracing subscriber on TestFramework init (#23)

* Set up tracing subscriber on TestFramework init

* Discard tracing subscriber already set errors
---
 Cargo.toml          |  2 ++
 src/batch_prover.rs |  4 ++--
 src/framework.rs    | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 353b79c..c80e3e1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,6 +13,7 @@ bollard = { version = "0.17.1" }
 futures = "0.3"
 hex = { version = "0.4.3", default-features = false, features = ["serde"] }
 jsonrpsee = { version = "0.24.2", features = ["http-client", "ws-client"] }
+log-panics = { version = "2", features = ["with-backtrace"] }
 rand = "0.8"
 serde = { version = "1.0.192", default-features = false, features = ["alloc", "derive"] }
 serde_json = { version = "1.0", default-features = false }
@@ -20,6 +21,7 @@ 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"] }
 
 # 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..5e66d87 100644
--- a/src/framework.rs
+++ b/src/framework.rs
@@ -1,7 +1,11 @@
-use std::{future::Future, sync::Arc};
+use std::{
+    future::Future,
+    sync::{Arc, Once},
+};
 
 use bitcoincore_rpc::RpcApi;
 use tracing::{debug, info};
+use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
 
 use super::{
     bitcoin::BitcoinNodeCluster, config::TestConfig, docker::DockerEnv, full_node::FullNode,
@@ -53,6 +57,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 +225,27 @@ impl TestFramework {
         Ok(())
     }
 }
+
+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() {
+            let _ = tracing_subscriber::registry()
+                .with(fmt::layer().json())
+                .with(env_filter)
+                .try_init();
+        } else {
+            let _ = tracing_subscriber::registry()
+                .with(fmt::layer())
+                .with(env_filter)
+                .try_init();
+        }
+
+        log_panics::init();
+    });
+}