/dd-java-agent/instrumentation/akka-http-10.0/src/test/scala/AkkaHttpTestAsyncWebServer.scala

https://github.com/DataDog/dd-trace-java · Scala · 64 lines · 58 code · 5 blank · 1 comment · 4 complexity · a992d121ed59314c0f134779ab4cfa54 MD5 · raw file

  1. import akka.actor.ActorSystem
  2. import akka.http.scaladsl.Http
  3. import akka.http.scaladsl.Http.ServerBinding
  4. import akka.http.scaladsl.model.HttpMethods.GET
  5. import akka.http.scaladsl.model._
  6. import akka.stream.ActorMaterializer
  7. import datadog.trace.agent.test.base.HttpServerTest
  8. import datadog.trace.agent.test.base.HttpServerTest.ServerEndpoint._
  9. import groovy.lang.Closure
  10. import scala.concurrent.{Await, ExecutionContextExecutor, Future}
  11. object AkkaHttpTestAsyncWebServer {
  12. implicit val system: ActorSystem = ActorSystem("my-system")
  13. implicit val materializer: ActorMaterializer = ActorMaterializer()
  14. // needed for the future flatMap/onComplete in the end
  15. implicit val executionContext: ExecutionContextExecutor = system.dispatcher
  16. val asyncHandler: HttpRequest => Future[HttpResponse] = {
  17. case HttpRequest(GET, uri: Uri, _, _, _) =>
  18. Future {
  19. val endpoint =
  20. HttpServerTest.ServerEndpoint.forPath(uri.path.toString())
  21. HttpServerTest.controller(
  22. endpoint,
  23. new Closure[HttpResponse](()) {
  24. def doCall(): HttpResponse = {
  25. val resp = HttpResponse(status = endpoint.getStatus) //.withHeaders(headers.Type)resp.contentType = "text/plain"
  26. endpoint match {
  27. case SUCCESS => resp.withEntity(endpoint.getBody)
  28. case QUERY_PARAM => resp.withEntity(uri.queryString().orNull)
  29. case REDIRECT =>
  30. resp.withHeaders(headers.Location(endpoint.getBody))
  31. case ERROR => resp.withEntity(endpoint.getBody)
  32. case EXCEPTION => throw new Exception(endpoint.getBody)
  33. case _ =>
  34. HttpResponse(status = NOT_FOUND.getStatus)
  35. .withEntity(NOT_FOUND.getBody)
  36. }
  37. }
  38. }
  39. )
  40. }
  41. }
  42. private var binding: ServerBinding = _
  43. def start(port: Int): Unit = synchronized {
  44. if (null == binding) {
  45. import scala.concurrent.duration._
  46. binding = Await.result(
  47. Http().bindAndHandleAsync(asyncHandler, "localhost", port),
  48. 10 seconds
  49. )
  50. }
  51. }
  52. def stop(): Unit = synchronized {
  53. if (null != binding) {
  54. binding.unbind()
  55. system.terminate()
  56. binding = null
  57. }
  58. }
  59. }