flashrom_tester: Only print color when stdout isatty

Add the atty crate as a dependency. Print log and report in color only
when isatty is true.

BUG=b:246250254
BRANCH=None
TEST=ssh dut flashrom_tester # no color
TEST=ssh -t dut flashrom_tester # color

Change-Id: Ia3cc527fb98e53eda6773622340cf10764df2cba
Signed-off-by: Evan Benn <evanbenn@chromium.org>
Reviewed-on: https://review.coreboot.org/c/flashrom/+/69270
Tested-by: build bot (Jenkins) <no-reply@coreboot.org>
Reviewed-by: Angel Pons <th3fanbus@gmail.com>
Reviewed-by: Edward O'Callaghan <quasisec@chromium.org>

4 files changed, 57 insertions, 25 deletions

diff --git a/util/flashrom_tester/Cargo.toml b/util/flashrom_tester/Cargo.toml
index b57f04e54..a7b688130 100644
--- a/util/flashrom_tester/Cargo.toml
+++ b/util/flashrom_tester/Cargo.toml
@@ -14,6 +14,7 @@ name = "flashrom_tester"
 required-features = ["cli"]
 
 [dependencies]
+atty = "0.2"
 built = { version = "0.5", features = ["chrono"] }
 chrono = { version = "0.4", optional = true }
 clap = { version = "2.33", default-features = false, optional = true }
diff --git a/util/flashrom_tester/src/logger.rs b/util/flashrom_tester/src/logger.rs
index 3c9b74b4e..fb26bde4c 100644
--- a/util/flashrom_tester/src/logger.rs
+++ b/util/flashrom_tester/src/logger.rs
@@ -41,6 +41,7 @@ use std::sync::Mutex;
 struct Logger<W: Write + Send> {
     level: log::LevelFilter,
     target: LogTarget<W>,
+    color: types::Color,
 }
 
 enum LogTarget<W>
@@ -57,16 +58,14 @@ impl<W: Write + Send> log::Log for Logger<W> {
     }
 
     fn log(&self, record: &log::Record) {
-        fn log_internal<W: Write>(mut w: W, record: &log::Record) -> std::io::Result<()> {
+        fn log_internal<W: Write>(
+            mut w: W,
+            record: &log::Record,
+            color: &types::Color,
+        ) -> std::io::Result<()> {
             let now = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Micros, true);
-            write!(w, "{}{} ", types::MAGENTA, now)?;
-            write!(
-                w,
-                "{}[ {} ]{} ",
-                types::YELLOW,
-                record.level(),
-                types::RESET
-            )?;
+            write!(w, "{}{} ", color.magenta, now)?;
+            write!(w, "{}[ {} ]{} ", color.yellow, record.level(), color.reset)?;
             writeln!(w, "{}", record.args())
         }
 
@@ -75,11 +74,11 @@ impl<W: Write + Send> log::Log for Logger<W> {
             LogTarget::Terminal => {
                 let stdout = std::io::stdout();
                 let mut lock = stdout.lock();
-                log_internal(&mut lock, record)
+                log_internal(&mut lock, record, &self.color)
             }
             LogTarget::Write(ref mutex) => {
                 let mut lock = mutex.lock().unwrap();
-                log_internal(&mut *lock, record)
+                log_internal(&mut *lock, record, &self.color)
             }
         };
     }
@@ -97,6 +96,11 @@ pub fn init(to_file: Option<PathBuf>, debug: bool) {
     let mut logger = Logger {
         level: log::LevelFilter::Info,
         target: LogTarget::Terminal,
+        color: if atty::is(atty::Stream::Stdout) {
+            types::COLOR
+        } else {
+            types::NOCOLOR
+        },
     };
 
     if debug {
@@ -106,6 +110,7 @@ pub fn init(to_file: Option<PathBuf>, debug: bool) {
         logger.target = LogTarget::Write(Mutex::new(
             std::fs::File::create(path).expect("Unable to open log file for writing"),
         ));
+        logger.color = types::NOCOLOR;
     }
 
     log::set_max_level(logger.level);
@@ -115,6 +120,7 @@ pub fn init(to_file: Option<PathBuf>, debug: bool) {
 #[cfg(test)]
 mod tests {
     use super::{LogTarget, Logger};
+    use flashrom_tester::types;
     use log::{Level, LevelFilter, Log, Record};
     use std::sync::Mutex;
 
@@ -125,6 +131,7 @@ mod tests {
             let logger = Logger {
                 level: LevelFilter::Info,
                 target: LogTarget::Write(lock),
+                color: types::COLOR,
             };
 
             for record in records {
diff --git a/util/flashrom_tester/src/tester.rs b/util/flashrom_tester/src/tester.rs
index 6ce3889e2..323b53656 100644
--- a/util/flashrom_tester/src/tester.rs
+++ b/util/flashrom_tester/src/tester.rs
@@ -573,6 +573,11 @@ pub fn collate_all_test_runs(
 ) {
     match format {
         OutputFormat::Pretty => {
+            let color = if atty::is(atty::Stream::Stdout) {
+                types::COLOR
+            } else {
+                types::NOCOLOR
+            };
             println!();
             println!("  =============================");
             println!("  =====  AVL qual RESULTS  ====");
@@ -591,8 +596,8 @@ pub fn collate_all_test_runs(
                 if *result != TestConclusion::Pass {
                     println!(
                         " {} {}",
-                        style!(format!(" <+> {} test:", name), types::BOLD),
-                        style_dbg!(result, types::RED)
+                        style!(format!(" <+> {} test:", name), color.bold, color),
+                        style_dbg!(result, color.red, color)
                     );
                     match error {
                         None => {}
@@ -601,8 +606,8 @@ pub fn collate_all_test_runs(
                 } else {
                     println!(
                         " {} {}",
-                        style!(format!(" <+> {} test:", name), types::BOLD),
-                        style_dbg!(result, types::GREEN)
+                        style!(format!(" <+> {} test:", name), color.bold, color),
+                        style_dbg!(result, color.green, color)
                     );
                 }
             }
diff --git a/util/flashrom_tester/src/types.rs b/util/flashrom_tester/src/types.rs
index b22ded2b4..9cefb27e9 100644
--- a/util/flashrom_tester/src/types.rs
+++ b/util/flashrom_tester/src/types.rs
@@ -33,21 +33,40 @@
 // Software Foundation.
 //
 
-pub const BOLD: &str = "\x1b[1m";
+pub struct Color {
+    pub bold: &'static str,
+    pub reset: &'static str,
+    pub magenta: &'static str,
+    pub yellow: &'static str,
+    pub green: &'static str,
+    pub red: &'static str,
+}
+
+pub const COLOR: Color = Color {
+    bold: "\x1b[1m",
+    reset: "\x1b[0m",
+    magenta: "\x1b[35m",
+    yellow: "\x1b[33m",
+    green: "\x1b[92m",
+    red: "\x1b[31m",
+};
 
-pub const RESET: &str = "\x1b[0m";
-pub const MAGENTA: &str = "\x1b[35m";
-pub const YELLOW: &str = "\x1b[33m";
-pub const GREEN: &str = "\x1b[92m";
-pub const RED: &str = "\x1b[31m";
+pub const NOCOLOR: Color = Color {
+    bold: "",
+    reset: "",
+    magenta: "",
+    yellow: "",
+    green: "",
+    red: "",
+};
 
 macro_rules! style_dbg {
-    ($s: expr, $c: expr) => {
-        format!("{}{:?}{}", $c, $s, types::RESET)
+    ($s: expr, $c: expr, $col: expr) => {
+        format!("{}{:?}{}", $c, $s, $col.reset)
     };
 }
 macro_rules! style {
-    ($s: expr, $c: expr) => {
-        format!("{}{}{}", $c, $s, types::RESET)
+    ($s: expr, $c: expr, $col: expr) => {
+        format!("{}{}{}", $c, $s, $col.reset)
     };
 }