From bd861fd25ef5b42c1329aea8efb50406b48824a9 Mon Sep 17 00:00:00 2001
From: Jane Lusby <jlusby42@gmail.com>
Date: Thu, 17 Sep 2020 13:20:30 -0700
Subject: [PATCH] update panic hook for zebra-test to supress confusing output
 (#1065)

* update panic hook for zebra-test to supress confusing output

* remove outdated comment
---
 Cargo.lock            |  1 +
 zebra-test/Cargo.toml |  1 +
 zebra-test/src/lib.rs | 54 ++++++++++++++++++++++++++++++++++++-------
 3 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7504e2fd..a5eb7a0d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3310,6 +3310,7 @@ dependencies = [
  "futures",
  "hex",
  "lazy_static",
+ "owo-colors",
  "pretty_assertions",
  "regex",
  "spandoc",
diff --git a/zebra-test/Cargo.toml b/zebra-test/Cargo.toml
index 212d2c81..7c426c22 100644
--- a/zebra-test/Cargo.toml
+++ b/zebra-test/Cargo.toml
@@ -20,6 +20,7 @@ spandoc = "0.2.0"
 regex = "1.3.9"
 thiserror = "1.0.20"
 pretty_assertions = "0.6.1"
+owo-colors = "1.1.3"
 
 [dev-dependencies]
 tokio = { version = "0.2", features = ["full"] }
diff --git a/zebra-test/src/lib.rs b/zebra-test/src/lib.rs
index 77455ca1..8c252320 100644
--- a/zebra-test/src/lib.rs
+++ b/zebra-test/src/lib.rs
@@ -1,8 +1,11 @@
 //! Miscellaneous test code for Zebra.
-use std::sync::Once;
+use color_eyre::section::PanicMessage;
+use owo_colors::OwoColorize;
 use tracing_error::ErrorLayer;
 use tracing_subscriber::{fmt, prelude::*, EnvFilter};
 
+use std::sync::Once;
+
 pub mod command;
 pub mod prelude;
 pub mod transcript;
@@ -70,14 +73,49 @@ pub fn init() {
                 });
             }))
             .display_env_section(false)
-            // Once I make a release with
-            // https://github.com/yaahc/color-eyre/pull/57 I'm going to add a
-            // custom PanicMessage handler for skipping the message in panics
-            // when the message is the exact one created by panics in tests that
-            // returned an Err. It will still print the content of the Err in
-            // the case where an error is returned.
-            // .panic_message(SkipTestReturnedErrPanicMessages)
+            .panic_message(SkipTestReturnedErrPanicMessages)
             .install()
             .unwrap();
     })
 }
+
+struct SkipTestReturnedErrPanicMessages;
+
+impl PanicMessage for SkipTestReturnedErrPanicMessages {
+    fn display(
+        &self,
+        pi: &std::panic::PanicInfo<'_>,
+        f: &mut std::fmt::Formatter<'_>,
+    ) -> std::fmt::Result {
+        let payload = pi
+            .payload()
+            .downcast_ref::<String>()
+            .map(String::as_str)
+            .or_else(|| pi.payload().downcast_ref::<&str>().cloned())
+            .unwrap_or("<non string panic payload>");
+
+        // skip panic output that is clearly from tests that returned an `Err`
+        // and assume that the test handler has already printed the value inside
+        // of the `Err`.
+        if payload.contains("the test returned a termination value with a non-zero status code") {
+            return write!(f, "---- end of test output ----");
+        }
+
+        writeln!(f, "{}", "\nThe application panicked (crashed).".red())?;
+
+        write!(f, "Message:  ")?;
+        writeln!(f, "{}", payload.cyan())?;
+
+        // If known, print panic location.
+        write!(f, "Location: ")?;
+        if let Some(loc) = pi.location() {
+            write!(f, "{}", loc.file().purple())?;
+            write!(f, ":")?;
+            write!(f, "{}", loc.line().purple())?;
+        } else {
+            write!(f, "<unknown>")?;
+        }
+
+        Ok(())
+    }
+}