From 602b67f20fb9486fe58f7f5ec8352dce7efb10a6 Mon Sep 17 00:00:00 2001
From: Andriy Plokhotnyuk <plokhotnyuk@gmail.com>
Date: Sat, 21 Dec 2024 11:07:05 +0100
Subject: [PATCH] More efficient cleaning of exception stack traces in zio-test

---
 .../src/main/scala/zio/test/TestArrow.scala   | 21 +++++++------------
 1 file changed, 7 insertions(+), 14 deletions(-)

diff --git a/test/shared/src/main/scala/zio/test/TestArrow.scala b/test/shared/src/main/scala/zio/test/TestArrow.scala
index f7f8359d750e..b6dc739417c5 100644
--- a/test/shared/src/main/scala/zio/test/TestArrow.scala
+++ b/test/shared/src/main/scala/zio/test/TestArrow.scala
@@ -228,20 +228,13 @@ object TestArrow {
     }
 
   private def attempt[A](expr: => TestTrace[A]): TestTrace[A] =
-    try {
-      expr
-    } catch {
-      case NonFatal(exception) =>
-        val trace = exception.getStackTrace
-        var met   = false
-        val newTrace = trace.filterNot { trace =>
-          if (trace.toString.contains("zio.test.TestArrow")) {
-            met = true
-          }
-          met
-        }
-        exception.setStackTrace(newTrace)
-        TestTrace.die(exception)
+    try expr
+    catch {
+      case ex if NonFatal(ex) =>
+        ex.setStackTrace(ex.getStackTrace.filterNot { (ste: StackTraceElement) =>
+          ste.getClassName.startsWith("zio.test.TestArrow")
+        })
+        TestTrace.die(ex)
     }
 
   def run[A, B](arrow: TestArrow[A, B], in: Either[Throwable, A]): TestTrace[B] = attempt {