1 files changed, 17 insertions, 23 deletions
diff --git a/kamon-core/src/test/scala/kamon/trace/ActiveSpanManagementSpec.scala b/kamon-core/src/test/scala/kamon/trace/ActiveSpanManagementSpec.scala
index ebee9f66..a6a7bc3a 100644
--- a/kamon-core/src/test/scala/kamon/trace/ActiveSpanManagementSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/ActiveSpanManagementSpec.scala
@@ -9,16 +9,16 @@ import org.scalatest.{Matchers, WordSpec}
 class ActiveSpanManagementSpec extends WordSpec with Matchers {
 
   "Kamon acting as a ActiveSpanSource" should {
-    "return a ActiveSpan wrapping a empty span when there is no currently active Span" in {
+    "return a empty span when there is no currently active Span" in {
       inspect(Kamon.activeSpan()) shouldBe empty
     }
 
-    "safely operate on a ActiveSpan that wraps a empty Span" in {
-      val activeSpan = Kamon.activeSpan()
+    "safely operate on a empty Span" in {
+      val emptySpan = Kamon.activeSpan()
       val activeSpanData = inspect(Kamon.activeSpan())
       activeSpanData shouldBe empty
 
-      activeSpan
+      emptySpan
         .setOperationName("test")
         .addBaggage("key", "value")
         .addMetricTag("key", "value")
@@ -28,39 +28,33 @@ class ActiveSpanManagementSpec extends WordSpec with Matchers {
         .addSpanTag("boolean-false", false)
         .annotate(Annotation(Clock.microTimestamp(), "event", Map("k" -> "v")))
 
-      val baggage = activeSpan.context().baggage
+      val baggage = emptySpan.context().baggage
       baggage.add("key", "value")
       baggage.get("key") shouldBe empty
       baggage.getAll() shouldBe empty
 
-      val continuation = activeSpan.capture()
-      val activatedSpan = continuation.activate()
-      inspect(Kamon.activeSpan()) shouldBe empty
-      activatedSpan.deactivate()
+      Kamon.withActiveSpan(emptySpan) {
+        inspect(Kamon.activeSpan()) shouldBe empty
+      }
 
       inspect(Kamon.activeSpan()) shouldBe empty
     }
 
-    "set a Span as active when using makeActive" in {
+    "set a Span as active when using activate" in {
       val span = Kamon.buildSpan("mySpan").start()
-      val activeSpan = Kamon.makeActive(span)
-      Kamon.activeSpan() shouldBe theSameInstanceAs(activeSpan)
-      activeSpan.deactivate()
-    }
-
-    "set a Span as active when using startActive" in {
-      val activeSpan = Kamon.buildSpan("mySpan").startActive()
-      Kamon.activeSpan() shouldBe theSameInstanceAs(activeSpan)
-      activeSpan.deactivate()
+      val scope = Kamon.activate(span)
+      Kamon.activeSpan() shouldBe theSameInstanceAs(span)
+      scope.close()
     }
 
-    "restore the previously active Span when a ActiveSpan gets deactivated" in {
+    "restore the previously active Span when a scope is closed" in {
       val previouslyActiveSpan = Kamon.activeSpan()
       inspect(Kamon.activeSpan()) shouldBe empty
 
-      val activeSpan = Kamon.buildSpan("mySpan").startActive()
-      Kamon.activeSpan() shouldBe theSameInstanceAs(activeSpan)
-      activeSpan.deactivate()
+      val span = Kamon.buildSpan("mySpan").start()
+      Kamon.withActiveSpan(span) {
+        Kamon.activeSpan() shouldBe theSameInstanceAs(span)
+      }
 
       Kamon.activeSpan() shouldBe theSameInstanceAs(previouslyActiveSpan)
     }