Code review improvements & gauges for in-progress requests

1 files changed, 118 insertions, 11 deletions

diff --git a/metrics/prometheus-backend/src/test/scala/com/softwaremill/sttp/prometheus/PrometheusBackendTest.scala b/metrics/prometheus-backend/src/test/scala/com/softwaremill/sttp/prometheus/PrometheusBackendTest.scala
index e6476de..33477a0 100644
--- a/metrics/prometheus-backend/src/test/scala/com/softwaremill/sttp/prometheus/PrometheusBackendTest.scala
+++ b/metrics/prometheus-backend/src/test/scala/com/softwaremill/sttp/prometheus/PrometheusBackendTest.scala
@@ -1,12 +1,22 @@
 package com.softwaremill.sttp.prometheus
 
-import com.softwaremill.sttp.{HttpURLConnectionBackend, Id, sttp}
-import io.prometheus.client.CollectorRegistry
-import org.scalatest.{BeforeAndAfter, FlatSpec, Matchers}
-import com.softwaremill.sttp._
+import java.lang
+import java.util.concurrent.CountDownLatch
+
 import com.softwaremill.sttp.testing.SttpBackendStub
+import com.softwaremill.sttp.{HttpURLConnectionBackend, Id, sttp, _}
+import io.prometheus.client.CollectorRegistry
+import org.scalatest.concurrent.{Eventually, IntegrationPatience}
+import org.scalatest.{BeforeAndAfter, FlatSpec, Matchers, OptionValues}
 
-class PrometheusBackendTest extends FlatSpec with Matchers with BeforeAndAfter {
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.Future
+
+class PrometheusBackendTest extends FlatSpec with Matchers with BeforeAndAfter with Eventually with OptionValues {
+
+  before {
+    CollectorRegistry.defaultRegistry.clear()
+  }
 
   it should "use default histogram name" in {
     // given
@@ -18,22 +28,119 @@ class PrometheusBackendTest extends FlatSpec with Matchers with BeforeAndAfter {
     (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
 
     // then
-    val result = CollectorRegistry.defaultRegistry.getSampleValue(s"${PrometheusBackend.DefaultHistogramName}_count")
-    result shouldBe requestsNumber
+    getMetricVale(s"${PrometheusBackend.DefaultHistogramName}_count").value shouldBe requestsNumber
   }
 
   it should "use mapped request to histogram name" in {
     // given
-    val customHistogramName = "my-custom-histogram"
+    val customHistogramName = "my_custom_histogram"
     val backend =
-      PrometheusBackend[Id, Nothing](SttpBackendStub(HttpURLConnectionBackend()), Some(_ => customHistogramName))
+      PrometheusBackend[Id, Nothing](SttpBackendStub(HttpURLConnectionBackend()), _ => Some(customHistogramName))
     val requestsNumber = 5
 
     // when
     (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
 
     // then
-    CollectorRegistry.defaultRegistry.getSampleValue(s"${PrometheusBackend.DefaultHistogramName}_count") shouldBe null
-    CollectorRegistry.defaultRegistry.getSampleValue(s"${customHistogramName}_count") shouldBe requestsNumber
+    getMetricVale(s"${PrometheusBackend.DefaultHistogramName}_count") shouldBe empty
+    getMetricVale(s"${customHistogramName}_count").value shouldBe requestsNumber
+  }
+
+  it should "disable histograms" in {
+    // given
+    val backend =
+      PrometheusBackend[Id, Nothing](SttpBackendStub(HttpURLConnectionBackend()), _ => None)
+    val requestsNumber = 6
+
+    // when
+    (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
+
+    // then
+    getMetricVale(s"${PrometheusBackend.DefaultHistogramName}_count") shouldBe empty
+  }
+
+  it should "use default gauge name" in {
+    // given
+    val requestsNumber = 10
+    val countDownLatch = new CountDownLatch(1)
+    val backendStub = SttpBackendStub.asynchronousFuture.whenAnyRequest.thenRespondWrapped {
+      Future {
+        countDownLatch.await()
+        Response(Right(""), 200, "", Nil, Nil)
+      }
+    }
+    val backend = PrometheusBackend[Future, Nothing](backendStub)
+
+    // when
+    (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
+
+    // then
+    eventually {
+      getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName).value shouldBe requestsNumber
+    }
+
+    countDownLatch.countDown()
+    eventually {
+      getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName).value shouldBe 0
+    }
+  }
+
+  it should "use mapped request to gauge name" in {
+    // given
+    val customGaugeName = "my_custom_gauge"
+    val requestsNumber = 10
+    val countDownLatch = new CountDownLatch(1)
+    val backendStub = SttpBackendStub.asynchronousFuture.whenAnyRequest.thenRespondWrapped {
+      Future {
+        countDownLatch.await()
+        Response(Right(""), 200, "", Nil, Nil)
+      }
+    }
+    val backend =
+      PrometheusBackend[Future, Nothing](backendStub, requestToInProgressGaugeNameMapper = _ => Some(customGaugeName))
+
+    // when
+    (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
+
+    // then
+    eventually {
+      getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName) shouldBe empty
+      getMetricVale(customGaugeName).value shouldBe requestsNumber
+    }
+
+    countDownLatch.countDown()
+    eventually {
+      getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName) shouldBe empty
+      getMetricVale(customGaugeName).value shouldBe 0
+    }
   }
+
+  it should "disable gauge" in {
+    // given
+    val requestsNumber = 10
+    val countDownLatch = new CountDownLatch(1)
+    val backendStub = SttpBackendStub.asynchronousFuture.whenAnyRequest.thenRespondWrapped {
+      Future {
+        countDownLatch.await()
+        Response(Right(""), 200, "", Nil, Nil)
+      }
+    }
+    val backend = PrometheusBackend[Future, Nothing](backendStub, requestToInProgressGaugeNameMapper = _ => None)
+
+    // when
+    (0 until requestsNumber).foreach(_ => backend.send(sttp.get(uri"http://127.0.0.1/foo")))
+
+    // then
+    getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName) shouldBe empty
+
+    countDownLatch.countDown()
+    eventually {
+      getMetricVale(s"${PrometheusBackend.DefaultHistogramName}_count").value shouldBe requestsNumber
+      getMetricVale(PrometheusBackend.DefaultRequestsInProgressGaugeName) shouldBe empty
+    }
+  }
+
+  private[this] def getMetricVale(name: String): Option[lang.Double] =
+    Option(CollectorRegistry.defaultRegistry.getSampleValue(name))
+
 }