Additional metric inspections

1 files changed, 22 insertions, 1 deletions

diff --git a/kamon-testkit/src/main/scala/kamon/testkit/MetricInspection.scala b/kamon-testkit/src/main/scala/kamon/testkit/MetricInspection.scala
index bb6d27f9..be340df9 100644
--- a/kamon-testkit/src/main/scala/kamon/testkit/MetricInspection.scala
+++ b/kamon-testkit/src/main/scala/kamon/testkit/MetricInspection.scala
@@ -15,7 +15,7 @@
 
 package kamon.testkit
 
-import kamon.metric._
+import kamon.metric.{BaseMetric, _}
 
 import _root_.scala.collection.concurrent.TrieMap
 
@@ -30,6 +30,27 @@ trait MetricInspection {
       val instrumentsWithTheTag = instruments.keys.filter(_.keys.exists(_ == tag))
       instrumentsWithTheTag.map(t => t(tag)).toSeq
     }
+
+    def partialRefine(tags: Map[String, String]): Seq[Map[String, String]] = {
+      val instrumentsField = classOf[BaseMetric[_, _]].getDeclaredField("instruments")
+      instrumentsField.setAccessible(true)
+
+      val instruments = instrumentsField.get(metric).asInstanceOf[TrieMap[Map[String, String], _]]
+
+      instruments.keys.filter { metricKey =>
+        tags.toSeq.forall { case (k, v) =>
+          metricKey.contains(k) && metricKey(k) == v
+        }
+      }.toSeq
+    }
+
+    def partialRefineKeys(tags: Set[String]): Seq[Map[String, String]] = {
+      val instrumentsField = classOf[BaseMetric[_, _]].getDeclaredField("instruments")
+      instrumentsField.setAccessible(true)
+
+      val instruments = instrumentsField.get(metric).asInstanceOf[TrieMap[Map[String, String], _]]
+      instruments.keys.filter(key => tags.subsetOf(key.keySet)).toSeq
+    }
   }
 
   implicit class HistogramMetricSyntax(histogram: Histogram) {