Mark error to Suite/Run when tests failed (#61)

* ConcurrentHashMap -> TrieMap * Tweak * Mark error to Suite/Run when tests failed
NomadBlacky · Jul 16, 2024 · 4105ae6 · 4105ae6
1 parent 270966c
commit 4105ae6
Showing 1 changed file with 37 additions and 21 deletions.
diff --git a/...er/src/main/scala/dev/nomadblacky/scalatest_otel_reporter/OpenTelemetryTestReporter.scala b/...er/src/main/scala/dev/nomadblacky/scalatest_otel_reporter/OpenTelemetryTestReporter.scala
@@ -6,8 +6,10 @@ import io.opentelemetry.context.Context
 import org.scalatest.Reporter
 import org.scalatest.events._
 
+import java.util.Collections
 import java.util.concurrent.ConcurrentHashMap
 import java.util.logging.Logger
+import scala.collection.concurrent.TrieMap
 
 trait OpenTelemetryTestReporter[A <: OpenTelemetry] extends Reporter {
   import OpenTelemetryTestReporter._
@@ -20,10 +22,12 @@ trait OpenTelemetryTestReporter[A <: OpenTelemetry] extends Reporter {
   private lazy val tracer = otel.getTracerProvider.get("scalatest")
   private lazy val logger = Logger.getLogger(classOf[OpenTelemetryTestReporter.type].getName)
 
-  // TODO: Extract spans to a container class
-  private var testRootSpan: Span = _
-  private val suitesMap          = new ConcurrentHashMap[String, Span]()
-  private val testsMap           = new ConcurrentHashMap[String, Span]()
+  private var testRootSpan: Option[Span] = None
+  private val suitesMap                  = TrieMap.empty[String, Span]
+  private val testsMap                   = TrieMap.empty[String, Span]
+
+  private val failedSuiteIds =
+    Collections.newSetFromMap[String](new ConcurrentHashMap[String, java.lang.Boolean])
 
   def apply(event: Event): Unit =
     event match {
@@ -33,25 +37,28 @@ trait OpenTelemetryTestReporter[A <: OpenTelemetry] extends Reporter {
       case starting: RunStarting =>
         logger.fine(s"RunStarting")
         val rootSpanName = starting.configMap.getWithDefault(ConfigKeyRootSpanName, DefaultRootSpanName)
-        testRootSpan = tracer.spanBuilder(rootSpanName).startSpan()
+        testRootSpan = Some(tracer.spanBuilder(rootSpanName).startSpan())
 
       case completed: RunCompleted =>
         logger.fine(s"RunCompleted")
-        Option(testRootSpan).fold(throw new IllegalStateException("TestRootSpan not found")) { span =>
-          span.end()
+        testRootSpan.fold(throw new IllegalStateException("TestRootSpan not found")) { span =>
+          val statusCode = if (failedSuiteIds.isEmpty) StatusCode.OK else StatusCode.ERROR
+          span
+            .setStatus(statusCode)
+            .end()
         }
         shutdownOtel()
 
       case stopped: RunStopped =>
         logger.fine(s"RunStopped")
-        Option(testRootSpan).fold(throw new IllegalStateException("TestRootSpan not found")) { span =>
+        testRootSpan.fold(throw new IllegalStateException("TestRootSpan not found")) { span =>
           span.end()
         }
         shutdownOtel()
 
       case aborted: RunAborted =>
         logger.fine(s"RunAborted")
-        Option(testRootSpan).fold(throw new IllegalStateException("TestRootSpan not found")) { span =>
+        testRootSpan.fold(throw new IllegalStateException("TestRootSpan not found")) { span =>
           span
             .setStatus(StatusCode.ERROR, aborted.message)
             .recordException(aborted.throwable.orNull)
@@ -64,22 +71,27 @@ trait OpenTelemetryTestReporter[A <: OpenTelemetry] extends Reporter {
        */
       case starting: SuiteStarting =>
         logger.fine(s"SuiteStarting: ${starting.suiteName}")
-        val suiteSpan =
-          tracer.spanBuilder(starting.suiteName).setParent(Context.current().`with`(testRootSpan)).startSpan()
-        suitesMap.put(starting.suiteId, suiteSpan)
+        testRootSpan.fold(throw new IllegalStateException("TestRootSpan not found")) { rootSpan =>
+          val suiteSpan =
+            tracer.spanBuilder(starting.suiteName).setParent(Context.current().`with`(rootSpan)).startSpan()
+          suitesMap.put(starting.suiteId, suiteSpan)
+        }
 
       case completed: SuiteCompleted =>
         logger.fine(s"SuiteCompleted: ${completed.suiteName}")
-        Option(suitesMap.remove(completed.suiteId))
+        suitesMap
+          .remove(completed.suiteId)
           .fold(throw new IllegalStateException(s"Suite not found: $completed")) { span =>
+            val statusCode = if (failedSuiteIds.contains(completed.suiteId)) StatusCode.ERROR else StatusCode.OK
             span
-              .setStatus(StatusCode.OK)
+              .setStatus(statusCode)
               .end()
           }
 
       case aborted: SuiteAborted =>
         logger.fine(s"SuiteAborted: ${aborted.suiteName}")
-        Option(suitesMap.remove(aborted.suiteId))
+        suitesMap
+          .remove(aborted.suiteId)
           .fold(throw new IllegalStateException(s"Suite not found: $aborted")) { span =>
             span
               .setStatus(StatusCode.ERROR, aborted.message)
@@ -92,14 +104,15 @@ trait OpenTelemetryTestReporter[A <: OpenTelemetry] extends Reporter {
        */
       case starting: TestStarting =>
         logger.fine(s"TestStarting: ${starting.testName}")
-        val suiteSpan     = Option(suitesMap.get(starting.suiteId))
+        val suiteSpan     = suitesMap.get(starting.suiteId)
         val parentContext = suiteSpan.fold(Context.current())(Context.current().`with`)
         val testSpan      = tracer.spanBuilder(starting.testName).setParent(parentContext).startSpan()
         testsMap.put(starting.testName, testSpan)
 
       case succeeded: TestSucceeded =>
         logger.fine(s"TestSucceeded: ${succeeded.testName}")
-        Option(testsMap.remove(succeeded.testName))
+        testsMap
+          .remove(succeeded.testName)
           .fold(throw new IllegalStateException(s"Test not found: $succeeded")) { span =>
             span
               .setStatus(StatusCode.OK)
@@ -108,8 +121,10 @@ trait OpenTelemetryTestReporter[A <: OpenTelemetry] extends Reporter {
 
       case failed: TestFailed =>
         logger.fine(s"TestFailed: ${failed.testName}")
-        Option(testsMap.remove(failed.testName))
+        testsMap
+          .remove(failed.testName)
           .fold(throw new IllegalStateException(s"Test not found: ${failed.testName}")) { span =>
+            failedSuiteIds.add(failed.suiteId)
             span
               .setStatus(StatusCode.ERROR, failed.message)
               .recordException(failed.throwable.orNull)
@@ -118,21 +133,22 @@ trait OpenTelemetryTestReporter[A <: OpenTelemetry] extends Reporter {
 
       case ignored: TestIgnored =>
         logger.fine(s"TestIgnored: ${ignored.testName}")
-        val suiteSpan     = Option(suitesMap.get(ignored.suiteId))
+        val suiteSpan     = suitesMap.get(ignored.suiteId)
         val parentContext = suiteSpan.fold(Context.current())(Context.current().`with`)
         val testSpan      = tracer.spanBuilder(ignored.testName).setParent(parentContext).startSpan()
         testSpan.end()
 
       case pending: TestPending =>
         logger.fine(s"TestPending: ${pending.testName}")
-        val suiteSpan     = Option(suitesMap.get(pending.suiteId))
+        val suiteSpan     = suitesMap.get(pending.suiteId)
         val parentContext = suiteSpan.fold(Context.current())(Context.current().`with`)
         val testSpan      = tracer.spanBuilder(pending.testName).setParent(parentContext).startSpan()
         testSpan.end()
 
       case canceled: TestCanceled =>
         logger.fine(s"TestCanceled: ${canceled.testName}")
-        Option(testsMap.remove(canceled.testName))
+        testsMap
+          .remove(canceled.testName)
           .fold(throw new IllegalStateException(s"Test not found: ${canceled.testName}")) { span =>
             span.recordException(canceled.throwable.orNull).end()
           }