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

sermant-io · Aug 8, 2023 · 4fed3f6 · 4fed3f6
1 parent 4e1632d
commit 4fed3f6
Show file tree

Hide file tree

Showing 11 changed files with 802 additions and 64 deletions.
diff --git a/sermant-plugins/sermant-tag-transmission/tag-transmission-plugin/pom.xml b/sermant-plugins/sermant-tag-transmission/tag-transmission-plugin/pom.xml
@@ -20,6 +20,8 @@
         <rocketmq-client.version>4.8.0</rocketmq-client.version>
         <javax-servlet-api.version>3.0.1</javax-servlet-api.version>
         <kafka-clients.version>2.7.0</kafka-clients.version>
+        <tomcat-embed-core.version>9.0.38</tomcat-embed-core.version>
+        <powermock.version>2.0.9</powermock.version>
     </properties>
 
     <dependencies>
@@ -52,6 +54,46 @@
             <version>${kafka-clients.version}</version>
             <scope>provided</scope>
         </dependency>
+        <!--测试依赖-->
+        <dependency>
+            <groupId>junit</groupId>
+            <artifactId>junit</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-inline</artifactId>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.tomcat.embed</groupId>
+            <artifactId>tomcat-embed-core</artifactId>
+            <version>${tomcat-embed-core.version}</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.powermock</groupId>
+            <artifactId>powermock-core</artifactId>
+            <version>${powermock.version}</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.powermock</groupId>
+            <artifactId>powermock-api-mockito2</artifactId>
+            <version>${powermock.version}</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.powermock</groupId>
+            <artifactId>powermock-module-junit4</artifactId>
+            <version>${powermock.version}</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

diff --git a/...in/java/com/huaweicloud/sermant/tag/transmission/interceptors/HttpServletInterceptor.java b/...in/java/com/huaweicloud/sermant/tag/transmission/interceptors/HttpServletInterceptor.java
@@ -17,10 +17,7 @@
 package com.huaweicloud.sermant.tag.transmission.interceptors;
 
 import com.huaweicloud.sermant.core.plugin.agent.entity.ExecuteContext;
-import com.huaweicloud.sermant.core.plugin.agent.interceptor.AbstractInterceptor;
-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 java.util.Collections;
 import java.util.Enumeration;
@@ -36,24 +33,15 @@
  * @author tangle
  * @since 2023-07-18
  */
-public class HttpServletInterceptor extends AbstractInterceptor {
+public class HttpServletInterceptor extends AbstractServerInterceptor {
     /**
      * 过滤一次处理过程中拦截器的多次调用
      */
     private static final ThreadLocal<Boolean> LOCK_MARK = new ThreadLocal<>();
 
-    private final TagTransmissionConfig tagTransmissionConfig;
-
-    /**
-     * 构造器
-     */
-    public HttpServletInterceptor() {
-        tagTransmissionConfig = PluginConfigManager.getPluginConfig(TagTransmissionConfig.class);
-    }
-
     @Override
-    public ExecuteContext before(ExecuteContext context) {
-        if (!tagTransmissionConfig.isEnabled() || LOCK_MARK.get() != null) {
+    public ExecuteContext doBefore(ExecuteContext context) {
+        if (LOCK_MARK.get() != null) {
             return context;
         }
         LOCK_MARK.set(Boolean.TRUE);
@@ -77,7 +65,7 @@ public ExecuteContext before(ExecuteContext context) {
     }
 
     @Override
-    public ExecuteContext after(ExecuteContext context) {
+    public ExecuteContext doAfter(ExecuteContext context) {
         TrafficUtils.removeTrafficTag();
         LOCK_MARK.remove();
         return context;

diff --git a/...va/com/huaweicloud/sermant/tag/transmission/interceptors/RocketmqConsumerInterceptor.java b/...va/com/huaweicloud/sermant/tag/transmission/interceptors/RocketmqConsumerInterceptor.java
@@ -17,10 +17,7 @@
 package com.huaweicloud.sermant.tag.transmission.interceptors;
 
 import com.huaweicloud.sermant.core.plugin.agent.entity.ExecuteContext;
-import com.huaweicloud.sermant.core.plugin.agent.interceptor.AbstractInterceptor;
-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.apache.rocketmq.common.message.Message;
 
@@ -35,7 +32,7 @@
  * @author tangle
  * @since 2023-07-19
  */
-public class RocketmqConsumerInterceptor extends AbstractInterceptor {
+public class RocketmqConsumerInterceptor extends AbstractServerInterceptor {
     /**
      * getBody拦截方法的所在类名
      */
@@ -46,20 +43,25 @@ public class RocketmqConsumerInterceptor extends AbstractInterceptor {
      */
     private static final String ROCKETMQ_FILER_PREFIX = "org.apache.rocketmq";
 
-    private final TagTransmissionConfig tagTransmissionConfig;
-
-    /**
-     * 构造器
-     */
-    public RocketmqConsumerInterceptor() {
-        tagTransmissionConfig = PluginConfigManager.getPluginConfig(TagTransmissionConfig.class);
-    }
-
     @Override
-    public ExecuteContext before(ExecuteContext context) {
-        if (!tagTransmissionConfig.isEnabled()) {
+    public ExecuteContext doBefore(ExecuteContext context) {
+        if (!isAvailable()) {
             return context;
         }
+        if (context.getObject() instanceof Message) {
+            Message message = (Message) context.getObject();
+            Map<String, List<String>> tag = this.getTagFromMessage(message);
+            TrafficUtils.updateTrafficTag(tag);
+        }
+        return context;
+    }
+
+    /**
+     * 判断当前调用栈是否匹配拦截点进入的要求
+     *
+     * @return boolean
+     */
+    private boolean isAvailable() {
         StackTraceElement[] stackTraceElements = Thread.currentThread().getStackTrace();
         int stackTraceIdxMax = stackTraceElements.length - 1;
         for (int i = 0; i < stackTraceElements.length; i++) {
@@ -68,15 +70,11 @@ public ExecuteContext before(ExecuteContext context) {
             }
             if (i == stackTraceIdxMax || stackTraceElements[i + 1].getClassName()
                     .startsWith(ROCKETMQ_FILER_PREFIX)) {
-                return context;
+                return false;
             }
+            break;
         }
-        if (context.getObject() instanceof Message) {
-            Message message = (Message) context.getObject();
-            Map<String, List<String>> tag = this.getTagFromMessage(message);
-            TrafficUtils.updateTrafficTag(tag);
-        }
-        return context;
+        return true;
     }
 
     /**
@@ -99,7 +97,7 @@ private Map<String, List<String>> getTagFromMessage(Message message) {
     }
 
     @Override
-    public ExecuteContext after(ExecuteContext context) {
+    public ExecuteContext doAfter(ExecuteContext context) {
         return context;
     }
 }
diff --git a/...va/com/huaweicloud/sermant/tag/transmission/interceptors/RocketmqProducerInterceptor.java b/...va/com/huaweicloud/sermant/tag/transmission/interceptors/RocketmqProducerInterceptor.java
@@ -17,11 +17,8 @@
 package com.huaweicloud.sermant.tag.transmission.interceptors;
 
 import com.huaweicloud.sermant.core.plugin.agent.entity.ExecuteContext;
-import com.huaweicloud.sermant.core.plugin.agent.interceptor.AbstractInterceptor;
-import com.huaweicloud.sermant.core.plugin.config.PluginConfigManager;
 import com.huaweicloud.sermant.core.utils.tag.TrafficTag;
 import com.huaweicloud.sermant.core.utils.tag.TrafficUtils;
-import com.huaweicloud.sermant.tag.transmission.config.TagTransmissionConfig;
 
 import org.apache.rocketmq.common.protocol.header.SendMessageRequestHeader;
 
@@ -34,7 +31,7 @@
  * @author tangle
  * @since 2023-07-20
  */
-public class RocketmqProducerInterceptor extends AbstractInterceptor {
+public class RocketmqProducerInterceptor extends AbstractClientInterceptor {
     /**
      * SendMessageRequestHeader在sendMessage方法中的参数下标
      */
@@ -50,30 +47,12 @@ public class RocketmqProducerInterceptor extends AbstractInterceptor {
      */
     private static final char SPLIT_MARK = 2;
 
-    private final TagTransmissionConfig tagTransmissionConfig;
-
-    /**
-     * 构造器
-     */
-    public RocketmqProducerInterceptor() {
-        tagTransmissionConfig = PluginConfigManager.getPluginConfig(TagTransmissionConfig.class);
-    }
-
     @Override
-    public ExecuteContext before(ExecuteContext context) {
-        if (!tagTransmissionConfig.isEnabled()) {
-            return context;
-        }
-
-        TrafficTag trafficTag = TrafficUtils.getTrafficTag();
-        if (trafficTag == null) {
-            return context;
-        }
-
+    public ExecuteContext doBefore(ExecuteContext context) {
         if (context.getArguments()[ARGUMENT_INDEX] instanceof SendMessageRequestHeader) {
             SendMessageRequestHeader header = (SendMessageRequestHeader) context.getArguments()[ARGUMENT_INDEX];
             String oldProperties = header.getProperties();
-            String newProperties = this.insertTags2Properties(oldProperties, trafficTag);
+            String newProperties = this.insertTags2Properties(oldProperties, TrafficUtils.getTrafficTag());
             header.setProperties(newProperties);
         }
         return context;
@@ -109,7 +88,7 @@ private String insertTags2Properties(String oldProperties, TrafficTag trafficTag
     }
 
     @Override
-    public ExecuteContext after(ExecuteContext context) {
+    public ExecuteContext doAfter(ExecuteContext context) {
         TrafficUtils.removeTrafficTag();
         return context;
     }

diff --git a/...-plugin/src/test/java/com/huaweicloud/sermant/tag/transmission/interceptors/BaseTest.java b/...-plugin/src/test/java/com/huaweicloud/sermant/tag/transmission/interceptors/BaseTest.java
@@ -0,0 +1,60 @@
+/*
+ * Copyright (C) 2023-2023 Huawei Technologies Co., Ltd. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.huaweicloud.sermant.tag.transmission.interceptors;
+
+import com.huaweicloud.sermant.core.plugin.config.PluginConfigManager;
+import com.huaweicloud.sermant.tag.transmission.config.TagTransmissionConfig;
+
+import org.junit.After;
+import org.junit.Before;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * 流量标签透传UT基础测试类
+ *
+ * @author tangle
+ * @since 2023-07-27
+ */
+public class BaseTest {
+    public final TagTransmissionConfig tagTransmissionConfig = new TagTransmissionConfig();
+
+    public MockedStatic<PluginConfigManager> pluginConfigManagerMockedStatic;
+
+    public BaseTest() {
+        pluginConfigManagerMockedStatic = Mockito.mockStatic(PluginConfigManager.class);
+        pluginConfigManagerMockedStatic.when(() -> PluginConfigManager.getPluginConfig(TagTransmissionConfig.class))
+                .thenReturn(tagTransmissionConfig);
+    }
+
+    @Before
+    public void before() {
+        tagTransmissionConfig.setEnabled(true);
+        List<String> tagKeys = new ArrayList<>();
+        tagKeys.add("id");
+        tagKeys.add("name");
+        tagTransmissionConfig.setTagKeys(tagKeys);
+    }
+
+    @After
+    public void after() {
+        pluginConfigManagerMockedStatic.close();
+    }
+}