add clear agent service rule

apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/boot/ServiceManagerTest.java

@@ -3,7 +3,10 @@ package org.skywalking.apm.agent.core.boot;
 import java.lang.reflect.Field;
 import java.util.HashMap;
 import java.util.List;
+import org.junit.Rule;
 import org.junit.Test;
+import org.skywalking.apm.agent.core.boot.BootService;
+import org.skywalking.apm.agent.core.boot.ServiceManager;
 import org.skywalking.apm.agent.core.context.ContextManager;
 import org.skywalking.apm.agent.core.context.IgnoredTracerContext;
 import org.skywalking.apm.agent.core.context.TracingContext;
@@ -14,6 +17,7 @@ import org.skywalking.apm.agent.core.remote.GRPCChannelListener;
 import org.skywalking.apm.agent.core.remote.GRPCChannelManager;
 import org.skywalking.apm.agent.core.remote.TraceSegmentServiceClient;
 import org.skywalking.apm.agent.core.sampling.SamplingService;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
 
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -21,9 +25,12 @@ import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 
 public class ServiceManagerTest {
+
+    @Rule
+    public AgentServiceRule agentServiceRule = new AgentServiceRule();
+
     @Test
     public void testServiceDependencies() throws Exception {
-        ServiceManager.INSTANCE.boot();
         HashMap<Class, BootService> registryService = getFieldValue(ServiceManager.INSTANCE, "bootedServices");
 
         assertThat(registryService.size(), is(7));

apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/ContextManagerTest.java

 package org.skywalking.apm.agent.core.context;
 
-import com.google.instrumentation.trace.Span;
 import com.google.protobuf.InvalidProtocolBufferException;
-import java.util.Date;
 import java.util.List;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.BeforeClass;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.skywalking.apm.agent.core.boot.ServiceManager;
@@ -20,10 +19,11 @@ import org.skywalking.apm.agent.core.context.trace.TraceSegment;
 import org.skywalking.apm.agent.core.context.trace.TraceSegmentRef;
 import org.skywalking.apm.agent.core.context.util.AbstractTracingSpanHelper;
 import org.skywalking.apm.agent.core.context.util.SegmentHelper;
-import org.skywalking.apm.agent.core.context.util.SegmentStorage;
-import org.skywalking.apm.agent.core.context.util.SegmentStoragePoint;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
+import org.skywalking.apm.agent.core.test.tools.SegmentStorage;
+import org.skywalking.apm.agent.core.test.tools.SegmentStoragePoint;
 import org.skywalking.apm.agent.core.context.util.TraceSegmentRefHelper;
-import org.skywalking.apm.agent.core.context.util.TracingSegmentRunner;
+import org.skywalking.apm.agent.core.test.tools.TracingSegmentRunner;
 import org.skywalking.apm.agent.core.dictionary.DictionaryUtil;
 import org.skywalking.apm.network.proto.KeyWithStringValue;
 import org.skywalking.apm.network.proto.LogMessage;
@@ -45,10 +45,8 @@ public class ContextManagerTest {
     @SegmentStoragePoint
     private SegmentStorage tracingData;
 
-    @BeforeClass
-    public static void setUpBeforeClass() {
-        ServiceManager.INSTANCE.boot();
-    }
+    @Rule
+    public AgentServiceRule agentServiceRule = new AgentServiceRule();
 
     @Before
     public void setUp() throws Exception {

apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/IgnoredTracerContextTest.java

@@ -3,6 +3,7 @@ package org.skywalking.apm.agent.core.context;
 import java.util.LinkedList;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.skywalking.apm.agent.core.boot.ServiceManager;
@@ -10,9 +11,10 @@ import org.skywalking.apm.agent.core.conf.Config;
 import org.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
 import org.skywalking.apm.agent.core.context.trace.AbstractSpan;
 import org.skywalking.apm.agent.core.context.trace.NoopSpan;
-import org.skywalking.apm.agent.core.context.util.SegmentStorage;
-import org.skywalking.apm.agent.core.context.util.SegmentStoragePoint;
-import org.skywalking.apm.agent.core.context.util.TracingSegmentRunner;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
+import org.skywalking.apm.agent.core.test.tools.SegmentStorage;
+import org.skywalking.apm.agent.core.test.tools.SegmentStoragePoint;
+import org.skywalking.apm.agent.core.test.tools.TracingSegmentRunner;
 
 import static junit.framework.TestCase.assertNull;
 import static org.hamcrest.CoreMatchers.is;
@@ -24,6 +26,9 @@ public class IgnoredTracerContextTest {
     @SegmentStoragePoint
     private SegmentStorage storage;
 
+    @Rule
+    public AgentServiceRule agentServiceRule = new AgentServiceRule();
+
     @Before
     public void setUp() throws Exception {
         RemoteDownstreamConfig.Agent.APPLICATION_ID = 1;
@@ -57,7 +62,6 @@ public class IgnoredTracerContextTest {
 
     @Test
     public void ignoredTraceContextWithExcludeOperationName() {
-        ServiceManager.INSTANCE.boot();
         AbstractSpan abstractSpan = ContextManager.createEntrySpan("test.js", null);
         ContextManager.stopSpan();
 
@@ -68,7 +72,6 @@ public class IgnoredTracerContextTest {
 
     @Test
     public void ignoredTraceContextWithEmptyOperationName() {
-        ServiceManager.INSTANCE.boot();
         ContextCarrier contextCarrier = new ContextCarrier();
         AbstractSpan abstractSpan = ContextManager.createExitSpan("", contextCarrier, "127.0.0.1:2181");
         ContextManager.stopSpan();

apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/DiscoveryRestServiceClientTest.java

@@ -8,6 +8,7 @@ import org.junit.Rule;
 import org.junit.Test;
 import org.skywalking.apm.agent.core.conf.Config;
 import org.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
 
 import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
 import static com.github.tomakehurst.wiremock.client.WireMock.get;
@@ -18,6 +19,9 @@ import static org.hamcrest.MatcherAssert.assertThat;
 
 public class DiscoveryRestServiceClientTest {
 
+    @Rule
+    public AgentServiceRule agentServiceRule = new AgentServiceRule();
+
     private DiscoveryRestServiceClient client;
     @Rule
     public WireMockRule wireMockRule = new WireMockRule(8089);

apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/GRPCChannelManagerTest.java

@@ -18,6 +18,7 @@ import org.powermock.core.classloader.annotations.PrepareForTest;
 import org.powermock.modules.junit4.PowerMockRunner;
 import org.powermock.reflect.Whitebox;
 import org.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
 
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.MatcherAssert.assertThat;

apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/TraceSegmentServiceClientTest.java

@@ -17,9 +17,10 @@ import org.skywalking.apm.agent.core.context.ContextManager;
 import org.skywalking.apm.agent.core.context.tag.Tags;
 import org.skywalking.apm.agent.core.context.trace.AbstractSpan;
 import org.skywalking.apm.agent.core.context.trace.SpanLayer;
-import org.skywalking.apm.agent.core.context.util.SegmentStorage;
-import org.skywalking.apm.agent.core.context.util.SegmentStoragePoint;
-import org.skywalking.apm.agent.core.context.util.TracingSegmentRunner;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
+import org.skywalking.apm.agent.core.test.tools.SegmentStorage;
+import org.skywalking.apm.agent.core.test.tools.SegmentStoragePoint;
+import org.skywalking.apm.agent.core.test.tools.TracingSegmentRunner;
 import org.skywalking.apm.network.proto.Downstream;
 import org.skywalking.apm.network.proto.SpanObject;
 import org.skywalking.apm.network.proto.SpanType;
@@ -34,6 +35,10 @@ import static org.mockito.Mockito.spy;
 
 @RunWith(TracingSegmentRunner.class)
 public class TraceSegmentServiceClientTest {
+
+    @Rule
+    public AgentServiceRule agentServiceRule = new AgentServiceRule();
+
     @Rule
     public GrpcServerRule grpcServerRule = new GrpcServerRule().directExecutor();
 
@@ -73,7 +78,6 @@ public class TraceSegmentServiceClientTest {
 
     @Before
     public void setUp() throws Throwable {
-        ServiceManager.INSTANCE.boot();
         Whitebox.setInternalState(ServiceManager.INSTANCE.findService(GRPCChannelManager.class), "reconnect", false);
         spy(serviceClient);
 

apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/AgentServiceRule.java

+package org.skywalking.apm.agent.core.test.tools;
+
+import java.util.HashMap;
+import java.util.LinkedList;
+import org.junit.rules.ExternalResource;
+import org.powermock.reflect.Whitebox;
+import org.skywalking.apm.agent.core.boot.BootService;
+import org.skywalking.apm.agent.core.boot.ServiceManager;
+import org.skywalking.apm.agent.core.context.IgnoredTracerContext;
+import org.skywalking.apm.agent.core.context.TracingContext;
+import org.skywalking.apm.agent.core.context.TracingContextListener;
+
+public class AgentServiceRule extends ExternalResource {
+
+    @Override
+    protected void after() {
+        super.after();
+        Whitebox.setInternalState(ServiceManager.INSTANCE, "bootedServices", new HashMap<Class, BootService>());
+        Whitebox.setInternalState(TracingContext.ListenerManager.class, "LISTENERS", new LinkedList<TracingContextListener>() );
+        Whitebox.setInternalState(IgnoredTracerContext.ListenerManager.class, "LISTENERS", new LinkedList<TracingContextListener>() );
+    }
+
+    @Override
+    protected void before() throws Throwable {
+        super.before();
+        ServiceManager.INSTANCE.boot();
+    }
+}

apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/SegmentStorage.java → apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/SegmentStorage.java

-package org.skywalking.apm.agent.core.context.util;
+package org.skywalking.apm.agent.core.test.tools;
 
 import java.util.LinkedList;
 import java.util.List;
@@ -9,7 +9,7 @@ public class SegmentStorage {
     private LinkedList<TraceSegment> traceSegments;
     private LinkedList<IgnoredTracerContext> ignoredTracerContexts;
 
-    SegmentStorage() {
+   public SegmentStorage() {
         traceSegments = new LinkedList<TraceSegment>();
         ignoredTracerContexts = new LinkedList<IgnoredTracerContext>();
     }

apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/SegmentStoragePoint.java → apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/SegmentStoragePoint.java

-package org.skywalking.apm.agent.core.context.util;
+package org.skywalking.apm.agent.core.test.tools;
 
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;

apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/TracingSegmentRunner.java → apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/TracingSegmentRunner.java

-package org.skywalking.apm.agent.core.context.util;
+package org.skywalking.apm.agent.core.test.tools;
 
 import java.lang.reflect.Field;
-import org.junit.runner.Description;
-import org.junit.runner.notification.Failure;
-import org.junit.runner.notification.RunNotifier;
 import org.junit.runners.BlockJUnit4ClassRunner;
+import org.junit.runners.model.FrameworkMethod;
 import org.junit.runners.model.InitializationError;
+import org.junit.runners.model.Statement;
 import org.skywalking.apm.agent.core.context.IgnoreTracerContextListener;
 import org.skywalking.apm.agent.core.context.IgnoredTracerContext;
 import org.skywalking.apm.agent.core.context.TracingContext;
@@ -36,51 +35,39 @@ public class TracingSegmentRunner extends BlockJUnit4ClassRunner {
         return targetObject;
     }
 
-    @Override
-    public void run(RunNotifier notifier) {
-        notifier.addListener(new RunListener());
-        super.run(notifier);
-    }
+    @Override protected Statement withAfters(FrameworkMethod method, Object target, final Statement statement) {
+        return new Statement() {
+            @Override public void evaluate() throws Throwable {
+                if (field != null) {
+                    try {
+                        tracingData = new SegmentStorage();
+                        field.set(targetObject, tracingData);
+                    } catch (IllegalAccessException e) {
+                    }
+                }
+                tracingContextListener = new TracingContextListener() {
+                    @Override
+                    public void afterFinished(TraceSegment traceSegment) {
+                        tracingData.addTraceSegment(traceSegment);
+                    }
+                };
 
-    class RunListener extends org.junit.runner.notification.RunListener {
-        @Override
-        public void testStarted(Description description) throws Exception {
-            if (field != null) {
+                ignoreTracerContextListener = new IgnoreTracerContextListener() {
+                    @Override
+                    public void afterFinished(IgnoredTracerContext tracerContext) {
+                        tracingData.addIgnoreTraceContext(tracerContext);
+                    }
+                };
+
+                TracingContext.ListenerManager.add(tracingContextListener);
+                IgnoredTracerContext.ListenerManager.add(ignoreTracerContextListener);
                 try {
-                    tracingData = new SegmentStorage();
-                    field.set(targetObject, tracingData);
-                } catch (IllegalAccessException e) {
+                    statement.evaluate();
+                } finally {
+                    TracingContext.ListenerManager.remove(tracingContextListener);
+                    IgnoredTracerContext.ListenerManager.remove(ignoreTracerContextListener);
                 }
             }
-            tracingContextListener = new TracingContextListener() {
-                @Override
-                public void afterFinished(TraceSegment traceSegment) {
-                    tracingData.addTraceSegment(traceSegment);
-                }
-            };
-
-            ignoreTracerContextListener = new IgnoreTracerContextListener() {
-                @Override
-                public void afterFinished(IgnoredTracerContext tracerContext) {
-                    tracingData.addIgnoreTraceContext(tracerContext);
-                }
-            };
-            TracingContext.ListenerManager.add(tracingContextListener);
-            IgnoredTracerContext.ListenerManager.add(ignoreTracerContextListener);
-            super.testStarted(description);
-        }
-
-        @Override
-        public void testFinished(Description description) throws Exception {
-            super.testFinished(description);
-            TracingContext.ListenerManager.remove(tracingContextListener);
-            IgnoredTracerContext.ListenerManager.remove(ignoreTracerContextListener);
-        }
-
-        @Override
-        public void testFailure(Failure failure) throws Exception {
-            super.testFailure(failure);
-            TracingContext.ListenerManager.remove(tracingContextListener);
-        }
+        };
     }
 }