[Feature][core] test adapt metrics output

2 files changed, 14 insertions, 3 deletions

diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookupTest.java
index 7a6701e..b5df7e0 100644
--- a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookupTest.java
+++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookupTest.java
@@ -4,7 +4,7 @@ import com.geedgenetworks.common.Event;
 import com.geedgenetworks.common.udf.UDFContext;
 import org.apache.flink.api.common.functions.RuntimeContext;
 import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import java.util.*;
@@ -23,8 +23,8 @@ public class IntelligenceIndicatorLookupTest {
 
     private static RuntimeContext runtimeContext;
 
-    @BeforeAll
-    static void setUp() {
+    @BeforeEach
+    void setUp() {
         runtimeContext = mockRuntimeContext();
 
         String content = "type,ip_addr_format,ip1,ip2,domain,tags\nIP,CIDR,116.178.65.0,25,ali.com,\"阿里1,云服务1\"\nDomain,CIDR,116.178.65.0,25,ali.com,\"阿里2,云服务2\"";
diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java
index 0178375..20defff 100644
--- a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java
+++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java
@@ -11,6 +11,9 @@ import com.geedgenetworks.core.udf.knowlegdebase.handler.AbstractMultipleKnowled
 import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.functions.RuntimeContext;
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.metrics.SimpleCounter;
+import org.apache.flink.runtime.metrics.groups.OperatorMetricGroup;
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
 
@@ -44,6 +47,10 @@ public class LookupTestUtils {
         RuntimeContext runtimeContext = Mockito.mock(RuntimeContext.class);
         ExecutionConfig executionConfig = Mockito.mock(ExecutionConfig.class);
         Mockito.when(runtimeContext.getExecutionConfig()).thenReturn(executionConfig);
+        MetricGroup metricGroup = Mockito.mock(OperatorMetricGroup.class);
+        Mockito.when(runtimeContext.getMetricGroup()).thenReturn(metricGroup);
+        Mockito.when(metricGroup.addGroup(Mockito.anyString())).thenReturn(metricGroup);
+        Mockito.when(metricGroup.counter(Mockito.anyString())).thenReturn(new SimpleCounter());
         Configuration configuration = new Configuration();
         CommonConfig commonConfig = new CommonConfig();
         KnowledgeBaseConfig knowledgeBaseConfig = new KnowledgeBaseConfig();
@@ -92,6 +99,10 @@ public class LookupTestUtils {
         RuntimeContext runtimeContext = Mockito.mock(RuntimeContext.class);
         ExecutionConfig executionConfig = Mockito.mock(ExecutionConfig.class);
         Mockito.when(runtimeContext.getExecutionConfig()).thenReturn(executionConfig);
+        MetricGroup metricGroup = Mockito.mock(OperatorMetricGroup.class);
+        Mockito.when(runtimeContext.getMetricGroup()).thenReturn(metricGroup);
+        Mockito.when(metricGroup.addGroup(Mockito.anyString())).thenReturn(metricGroup);
+        Mockito.when(metricGroup.counter(Mockito.anyString())).thenReturn(new SimpleCounter());
         Configuration configuration = new Configuration();
         CommonConfig commonConfig = new CommonConfig();
         KnowledgeBaseConfig knowledgeBaseConfig = new KnowledgeBaseConfig();