【UT】增加流量标签透传的UT、完善流量标签透传拦截器继承结构

sermant-io · Aug 11, 2023 · 4497e5b · 4497e5b
1 parent 7368744
commit 4497e5b
Show file tree

Hide file tree

Showing 10 changed files with 48 additions and 26 deletions.
diff --git a/...va/com/huaweicloud/sermant/tag/transmission/interceptors/RocketmqConsumerInterceptor.java b/...va/com/huaweicloud/sermant/tag/transmission/interceptors/RocketmqConsumerInterceptor.java
@@ -46,7 +46,8 @@ public class RocketmqConsumerInterceptor extends AbstractServerInterceptor {
 
     @Override
     public ExecuteContext doBefore(ExecuteContext context) {
-        if (!isAvailable()) {
+        StackTraceElement[] stackTraceElements = Thread.currentThread().getStackTrace();
+        if (!isRocketMqStackTrace(stackTraceElements)) {
             return context;
         }
         if (context.getObject() instanceof Message) {
@@ -64,8 +65,7 @@ public ExecuteContext doBefore(ExecuteContext context) {
      *
      * @return boolean
      */
-    private boolean isAvailable() {
-        StackTraceElement[] stackTraceElements = Thread.currentThread().getStackTrace();
+    private boolean isRocketMqStackTrace(StackTraceElement[] stackTraceElements) {
         int stackTraceIdxMax = stackTraceElements.length - 1;
         for (int i = 0; i < stackTraceElements.length; i++) {
             if (!ROCKETMQ_SELECT_CLASSNAME.equals(stackTraceElements[i].getClassName())) {

diff --git a/...va/com/huaweicloud/sermant/tag/transmission/interceptors/RocketmqProducerInterceptor.java b/...va/com/huaweicloud/sermant/tag/transmission/interceptors/RocketmqProducerInterceptor.java
@@ -89,13 +89,6 @@ private String insertTags2Properties(String oldProperties, TrafficTag trafficTag
 
     @Override
     public ExecuteContext doAfter(ExecuteContext context) {
-        TrafficUtils.removeTrafficTag();
-        return context;
-    }
-
-    @Override
-    public ExecuteContext onThrow(ExecuteContext context) {
-        TrafficUtils.removeTrafficTag();
         return context;
     }
 }
diff --git a/...-transmission-plugin/src/test/java/com/huaweicloud/sermant/tag/transmission/BaseTest.java b/...-transmission-plugin/src/test/java/com/huaweicloud/sermant/tag/transmission/BaseTest.java
@@ -17,6 +17,7 @@
 package com.huaweicloud.sermant.tag.transmission;
 
 import com.huaweicloud.sermant.core.plugin.config.PluginConfigManager;
+import com.huaweicloud.sermant.core.utils.tag.TrafficUtils;
 import com.huaweicloud.sermant.tag.transmission.config.CrossThreadConfig;
 
 import org.junit.AfterClass;
@@ -38,6 +39,7 @@ public static void before() {
         mockPluginConfigManager = Mockito.mockStatic(PluginConfigManager.class);
         mockPluginConfigManager.when(() -> PluginConfigManager.getPluginConfig(CrossThreadConfig.class))
                 .thenReturn(new CrossThreadConfig());
+        TrafficUtils.removeTrafficTag();
     }
 
     /**

diff --git a/...g/transmission/interceptors/BaseTest.java → ...ion/interceptors/BaseInterceptorTest.java b/...g/transmission/interceptors/BaseTest.java → ...ion/interceptors/BaseInterceptorTest.java
@@ -17,6 +17,7 @@
 package com.huaweicloud.sermant.tag.transmission.interceptors;
 
 import com.huaweicloud.sermant.core.plugin.config.PluginConfigManager;
+import com.huaweicloud.sermant.core.utils.tag.TrafficUtils;
 import com.huaweicloud.sermant.tag.transmission.config.TagTransmissionConfig;
 
 import org.junit.After;
@@ -33,12 +34,12 @@
  * @author tangle
  * @since 2023-07-27
  */
-public class BaseTest {
+public class BaseInterceptorTest {
     public final TagTransmissionConfig tagTransmissionConfig = new TagTransmissionConfig();
 
     public MockedStatic<PluginConfigManager> pluginConfigManagerMockedStatic;
 
-    public BaseTest() {
+    public BaseInterceptorTest() {
         pluginConfigManagerMockedStatic = Mockito.mockStatic(PluginConfigManager.class);
         pluginConfigManagerMockedStatic.when(() -> PluginConfigManager.getPluginConfig(TagTransmissionConfig.class))
                 .thenReturn(tagTransmissionConfig);
@@ -51,6 +52,7 @@ public void before() {
         tagKeys.add("id");
         tagKeys.add("name");
         tagTransmissionConfig.setTagKeys(tagKeys);
+        TrafficUtils.removeTrafficTag();
     }
 
     @After

diff --git a/...rceptors/HttpClient4xInterceptorTest.java → ...tpClient4XInterceptorInterceptorTest.java b/...rceptors/HttpClient4xInterceptorTest.java → ...tpClient4XInterceptorInterceptorTest.java
@@ -36,12 +36,12 @@
  * @author tangle
  * @since 2023-07-27
  */
-public class HttpClient4xInterceptorTest extends BaseTest {
+public class HttpClient4XInterceptorInterceptorTest extends BaseInterceptorTest {
     private final HttpClient4xInterceptor interceptor;
 
     private final Object[] arguments;
 
-    public HttpClient4xInterceptorTest() {
+    public HttpClient4XInterceptorInterceptorTest() {
         interceptor = new HttpClient4xInterceptor();
         arguments = new Object[2];
     }

diff --git a/...erceptors/HttpServletInterceptorTest.java → ...ttpServletInterceptorInterceptorTest.java b/...erceptors/HttpServletInterceptorTest.java → ...ttpServletInterceptorInterceptorTest.java
@@ -28,7 +28,6 @@
 
 import java.lang.reflect.Field;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -39,12 +38,12 @@
  * @author tangle
  * @since 2023-07-27
  */
-public class HttpServletInterceptorTest extends BaseTest {
+public class HttpServletInterceptorInterceptorTest extends BaseInterceptorTest {
     private final HttpServletInterceptor interceptor;
 
     private final Object[] arguments;
 
-    public HttpServletInterceptorTest() {
+    public HttpServletInterceptorInterceptorTest() {
         interceptor = new HttpServletInterceptor();
         arguments = new Object[2];
     }

diff --git a/...s/KafkaConsumerRecordInterceptorTest.java → ...umerRecordInterceptorInterceptorTest.java b/...s/KafkaConsumerRecordInterceptorTest.java → ...umerRecordInterceptorInterceptorTest.java
@@ -36,10 +36,10 @@
  * @author tangle
  * @since 2023-07-27
  */
-public class KafkaConsumerRecordInterceptorTest extends BaseTest {
+public class KafkaConsumerRecordInterceptorInterceptorTest extends BaseInterceptorTest {
     private final KafkaConsumerRecordInterceptor interceptor;
 
-    public KafkaConsumerRecordInterceptorTest() {
+    public KafkaConsumerRecordInterceptorInterceptorTest() {
         interceptor = new KafkaConsumerRecordInterceptor();
     }
 

diff --git a/...ceptors/KafkaProducerInterceptorTest.java → ...kaProducerInterceptorInterceptorTest.java b/...ceptors/KafkaProducerInterceptorTest.java → ...kaProducerInterceptorInterceptorTest.java
@@ -39,12 +39,12 @@
  * @author tangle
  * @since 2023-07-27
  */
-public class KafkaProducerInterceptorTest extends BaseTest {
+public class KafkaProducerInterceptorInterceptorTest extends BaseInterceptorTest {
     private final KafkaProducerInterceptor interceptor;
 
     private final Object[] arguments;
 
-    public KafkaProducerInterceptorTest() {
+    public KafkaProducerInterceptorInterceptorTest() {
         interceptor = new KafkaProducerInterceptor();
         arguments = new Object[2];
     }

diff --git a/...tors/RocketmqConsumerInterceptorTest.java → ...mqConsumerInterceptorInterceptorTest.java b/...tors/RocketmqConsumerInterceptorTest.java → ...mqConsumerInterceptorInterceptorTest.java
@@ -23,9 +23,11 @@
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mockito;
 import org.powermock.api.mockito.PowerMockito;
 import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
+import org.powermock.reflect.Whitebox;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -41,13 +43,37 @@
  */
 @RunWith(PowerMockRunner.class)
 @PrepareForTest(RocketmqConsumerInterceptor.class)
-public class RocketmqConsumerInterceptorTest extends BaseTest {
+public class RocketmqConsumerInterceptorInterceptorTest extends BaseInterceptorTest {
     private final RocketmqConsumerInterceptor interceptor;
 
-    public RocketmqConsumerInterceptorTest() throws Exception {
+    public RocketmqConsumerInterceptorInterceptorTest() throws Exception {
         RocketmqConsumerInterceptor interceptorBase = new RocketmqConsumerInterceptor();
         interceptor = PowerMockito.spy(interceptorBase);
-        PowerMockito.when(interceptor, "isAvailable").thenReturn(true);
+        PowerMockito.doReturn(true).when(interceptor, "isRocketMqStackTrace", Mockito.any());
+    }
+
+    @Test
+    public void testIsAvailable() throws Exception {
+        StackTraceElement[] stackTraceElements = new StackTraceElement[2];
+        RocketmqConsumerInterceptor interceptorBase = new RocketmqConsumerInterceptor();
+
+        // 测试可通过
+        stackTraceElements[0] = new StackTraceElement("org.apache.rocketmq.common.message.Message", "xxxMethod",
+                "xxxFile", 1);
+        stackTraceElements[1] = new StackTraceElement("org.apache.xxx", "xxxMethod",
+                "xxxFile", 1);
+        boolean resultAccess = Whitebox.invokeMethod(interceptorBase, "isRocketMqStackTrace",
+                (Object) stackTraceElements);
+        Assert.assertTrue(resultAccess);
+
+        // 测试不可通过
+        stackTraceElements[0] = new StackTraceElement("org.apache.rocketmq.common.message.Message", "xxxMethod",
+                "xxxFile", 1);
+        stackTraceElements[1] = new StackTraceElement("org.apache.rocketmq", "xxxMethod",
+                "xxxFile", 1);
+        boolean resultRefuse = Whitebox.invokeMethod(interceptorBase, "isRocketMqStackTrace",
+                (Object) stackTraceElements);
+        Assert.assertFalse(resultRefuse);
     }
 
     @Test

diff --git a/...tors/RocketmqProducerInterceptorTest.java → ...mqProducerInterceptorInterceptorTest.java b/...tors/RocketmqProducerInterceptorTest.java → ...mqProducerInterceptorInterceptorTest.java
@@ -34,12 +34,12 @@
  * @author tangle
  * @since 2023-07-27
  */
-public class RocketmqProducerInterceptorTest extends BaseTest {
+public class RocketmqProducerInterceptorInterceptorTest extends BaseInterceptorTest {
     private final RocketmqProducerInterceptor interceptor;
 
     private final Object[] arguments;
 
-    public RocketmqProducerInterceptorTest() {
+    public RocketmqProducerInterceptorInterceptorTest() {
         interceptor = new RocketmqProducerInterceptor();
         arguments = new Object[12];
     }