From 015ece40418f9cc75c2774c5e4a1c55f1169f9e1 Mon Sep 17 00:00:00 2001
From: harryteng9527 <zhimao.teng@gmail.com>
Date: Sun, 19 Mar 2023 22:12:09 +0800
Subject: [PATCH] Fix MetricSensor test

---
 .../metrics/collector/MetricSensorTest.java   | 23 ++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/common/src/test/java/org/astraea/common/metrics/collector/MetricSensorTest.java b/common/src/test/java/org/astraea/common/metrics/collector/MetricSensorTest.java
index 12a926e073..123239caea 100644
--- a/common/src/test/java/org/astraea/common/metrics/collector/MetricSensorTest.java
+++ b/common/src/test/java/org/astraea/common/metrics/collector/MetricSensorTest.java
@@ -57,7 +57,15 @@ void testNoSwallowException() {
         (client, ignored) -> {
           throw new RuntimeException("xxx");
         };
-    var sensor = MetricSensor.of(List.of(badMetricSensor, goodMetricSensor)).get();
+    var sensor =
+        MetricSensor.of(
+                List.of(badMetricSensor, goodMetricSensor),
+                e -> {
+                  if (e instanceof RuntimeException) {
+                    throw new RuntimeException();
+                  }
+                })
+            .get();
     Assertions.assertThrows(
         RuntimeException.class,
         () -> sensor.fetch(Mockito.mock(MBeanClient.class), ClusterBean.EMPTY));
@@ -82,11 +90,20 @@ void testSensorsWithExceptionHandler() {
     Assertions.assertEquals(
         1, sensor.fetch(Mockito.mock(MBeanClient.class), ClusterBean.EMPTY).size());
 
-    Assertions.assertThrows(
-        RuntimeException.class,
+    Assertions.assertDoesNotThrow(
         () ->
             MetricSensor.of(List.of(metricSensor0, metricSensor2))
                 .get()
                 .fetch(Mockito.mock(MBeanClient.class), ClusterBean.EMPTY));
+    Assertions.assertThrows(
+        NoSuchElementException.class,
+        () ->
+            MetricSensor.of(
+                    List.of(metricSensor1, metricSensor2),
+                    e -> {
+                      if (e instanceof NoSuchElementException) throw new NoSuchElementException();
+                    })
+                .get()
+                .fetch(Mockito.mock(MBeanClient.class), ClusterBean.EMPTY));
   }
 }