switched back to scalatest and created simple instrumentation for futures body

2 files changed, 52 insertions, 9 deletions

diff --git a/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala b/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
index 2eb8d07a..97e81b44 100644
--- a/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
+++ b/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
@@ -1,25 +1,36 @@
 package kamon.instrumentation
 
-import scala.concurrent.{Await, Future}
-import org.specs2.mutable.Specification
+import scala.concurrent.Future
 import scala.concurrent.ExecutionContext.Implicits.global
-import scala.concurrent.duration.{FiniteDuration, DurationLong}
-import org.specs2.time.{ Duration => SpecsDuration }
+import org.scalatest.WordSpec
+import org.scalatest.matchers.MustMatchers
+import org.scalatest.concurrent.PatienceConfiguration
+import kamon.TraceContext
+import java.util.UUID
 
 
-class FutureInstrumentationSpec extends Specification {
-  import Await.result
-  implicit def specsDuration2Akka(duration: SpecsDuration): FiniteDuration = new DurationLong(duration.inMillis).millis
+class FutureInstrumentationSpec extends WordSpec with MustMatchers with ScalaFutures with PatienceConfiguration {
 
   "a instrumented Future" should {
     "preserve the transaction context available during the future creation" in {
+      new ContextAwareTest {
+        val future = Future { TraceContext.current.get }
 
+        whenReady(future) { result =>
+          result must be === context
+        }
+      }
     }
 
     "use the same context available at creation when executing the onComplete callback" in {
-      val future = Future { "hello" }
 
-      result(future, 100 millis) === "hello"
     }
   }
+
+  trait ContextAwareTest {
+    val context = TraceContext(UUID.randomUUID(), Nil)
+    TraceContext.set(context)
+  }
 }
+
+
diff --git a/src/test/scala/kamon/instrumentation/ScalaFutures.scala b/src/test/scala/kamon/instrumentation/ScalaFutures.scala
new file mode 100644
index 00000000..169b709c
--- /dev/null
+++ b/src/test/scala/kamon/instrumentation/ScalaFutures.scala
@@ -0,0 +1,32 @@
+package kamon.instrumentation
+
+import scala.concurrent.{Await, Future}
+import scala.concurrent.duration.Duration
+import scala.util.{Failure, Success}
+import org.scalatest.concurrent.Futures
+import java.util.concurrent.TimeUnit
+
+trait ScalaFutures extends Futures {
+  implicit def scalaFutureToFutureConcept[T](future: Future[T]): FutureConcept[T] = new FutureConcept[T] {
+    def eitherValue: Option[Either[Throwable, T]] = {
+      if(!future.isCompleted)
+        None
+      else
+        future.value match {
+          case None => None
+          case Some(t) => t match {
+            case Success(v) => Some(Right(v))
+            case Failure(e) => Some(Left(e))
+          }
+        }
+    }
+
+    def isExpired: Boolean = false   // Scala futures cant expire
+
+    def isCanceled: Boolean = false  // Scala futures cannot be cancelled
+
+    override def futureValue(implicit config: PatienceConfig): T = {
+      Await.result(future, Duration(config.timeout.totalNanos, TimeUnit.NANOSECONDS))
+    }
+  }
+}