src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

package kamon.instrumentation

import scala.concurrent.{Promise, Future}
import scala.concurrent.ExecutionContext.Implicits.global
import org.scalatest.{OptionValues, WordSpec}
import org.scalatest.matchers.MustMatchers
import org.scalatest.concurrent.PatienceConfiguration
import kamon.TraceContext
import java.util.UUID
import scala.util.Success


class FutureInstrumentationSpec extends WordSpec with MustMatchers with ScalaFutures with PatienceConfiguration with OptionValues {

  "a instrumented Future" when {
    "created in a thread that does have a TraceContext" must {
      "preserve the TraceContext" which {
        "should be available during the body's execution" in { new FutureWithContext {

            whenReady(futureWithContext) { result =>
              result.value must be === testContext
            }
          }
        }

        "should be available during the execution of onComplete callbacks" in { new FutureWithContext {
            val onCompleteContext = Promise[TraceContext]()

            futureWithContext.onComplete({
              case _ => onCompleteContext.complete(Success(TraceContext.current.get))
            })

            whenReady(onCompleteContext.future) { result =>
              result must be === testContext
            }
          }
        }
      }
    }

    "created in a thread that doest have a TraceContext" must {
      "not capture any TraceContext" in { new FutureWithoutContext{

          whenReady(futureWithoutContext) { result =>
            result must be === None
          }
        }
      }
    }
  }


  trait FutureWithContext {
    val testContext = TraceContext(UUID.randomUUID(), Nil)
    TraceContext.set(testContext)

    val futureWithContext = Future { TraceContext.current }
  }

  trait FutureWithoutContext {
    TraceContext.clear
    val futureWithoutContext = Future { TraceContext.current }
  }
}