1+ /*
2+ * Licensed to the Apache Software Foundation (ASF) under one or more
3+ * contributor license agreements. See the NOTICE file distributed with
4+ * this work for additional information regarding copyright ownership.
5+ * The ASF licenses this file to You under the Apache License, Version 2.0
6+ * (the "License"); you may not use this file except in compliance with
7+ * the License. You may obtain a copy of the License at
8+ *
9+ * http://www.apache.org/licenses/LICENSE-2.0
10+ *
11+ * Unless required by applicable law or agreed to in writing, software
12+ * distributed under the License is distributed on an "AS IS" BASIS,
13+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+ * See the License for the specific language governing permissions and
15+ * limitations under the License.
16+ *
17+ */
18+
19+ package org .apache .skywalking .apm .plugin .sofarpc ;
20+
21+ import com .alipay .remoting .InvokeCallback ;
22+ import java .util .List ;
23+ import java .util .concurrent .CountDownLatch ;
24+ import java .util .concurrent .Executor ;
25+ import java .util .concurrent .Executors ;
26+ import org .apache .skywalking .apm .agent .core .context .ContextManager ;
27+ import org .apache .skywalking .apm .agent .core .context .trace .AbstractTracingSpan ;
28+ import org .apache .skywalking .apm .agent .core .context .trace .TraceSegment ;
29+ import org .apache .skywalking .apm .agent .test .helper .SegmentHelper ;
30+ import org .apache .skywalking .apm .agent .test .helper .SpanHelper ;
31+ import org .apache .skywalking .apm .agent .test .tools .AgentServiceRule ;
32+ import org .apache .skywalking .apm .agent .test .tools .SegmentStorage ;
33+ import org .apache .skywalking .apm .agent .test .tools .SegmentStoragePoint ;
34+ import org .apache .skywalking .apm .agent .test .tools .TracingSegmentRunner ;
35+ import org .junit .Assert ;
36+ import org .junit .Before ;
37+ import org .junit .Rule ;
38+ import org .junit .Test ;
39+ import org .junit .runner .RunWith ;
40+ import org .mockito .junit .MockitoJUnit ;
41+ import org .mockito .junit .MockitoRule ;
42+
43+ import static org .hamcrest .CoreMatchers .is ;
44+ import static org .junit .Assert .assertEquals ;
45+ import static org .junit .Assert .assertThat ;
46+
47+ @ RunWith (TracingSegmentRunner .class )
48+ public class InvokeCallbackWrapperTest {
49+
50+ @ SegmentStoragePoint
51+ private SegmentStorage segmentStorage ;
52+
53+ private Executor executor = Executors .newFixedThreadPool (1 );
54+
55+ @ Rule
56+ public AgentServiceRule agentServiceRule = new AgentServiceRule ();
57+ @ Rule
58+ public MockitoRule rule = MockitoJUnit .rule ();
59+
60+ private InvokeCallback callback ;
61+
62+ @ Before
63+ public void before () {
64+ callback = new InvokeCallback () {
65+ @ Override
66+ public void onResponse (final Object o ) {
67+ }
68+
69+ @ Override
70+ public void onException (final Throwable throwable ) {
71+ }
72+
73+ @ Override
74+ public Executor getExecutor () {
75+ return null ;
76+ }
77+ };
78+ }
79+
80+ static class WrapperWrapper implements InvokeCallback {
81+
82+ private InvokeCallback callback ;
83+
84+ private CountDownLatch countDownLatch ;
85+
86+ public CountDownLatch getCountDownLatch () {
87+ return countDownLatch ;
88+ }
89+
90+ public WrapperWrapper (InvokeCallback callback ) {
91+ this .countDownLatch = new CountDownLatch (1 );
92+ this .callback = callback ;
93+ }
94+
95+ @ Override
96+ public void onResponse (final Object o ) {
97+ callback .onResponse (o );
98+ countDownLatch .countDown ();
99+ }
100+
101+ @ Override
102+ public void onException (final Throwable throwable ) {
103+ callback .onException (throwable );
104+ countDownLatch .countDown ();
105+ }
106+
107+ @ Override
108+ public Executor getExecutor () {
109+ return null ;
110+ }
111+ }
112+
113+ @ Test
114+ public void testConstruct () {
115+ InvokeCallbackWrapper wrapper = new InvokeCallbackWrapper (callback );
116+ Assert .assertSame (callback , wrapper .getInvokeCallback ());
117+ Assert .assertNull (wrapper .getContextSnapshot ());
118+
119+ ContextManager .createEntrySpan ("sofarpc" , null );
120+ wrapper = new InvokeCallbackWrapper (callback );
121+ Assert .assertSame (callback , wrapper .getInvokeCallback ());
122+ Assert .assertEquals (ContextManager .getGlobalTraceId (), wrapper .getContextSnapshot ().getTraceId ().getId ());
123+ Assert .assertEquals ("sofarpc" , wrapper .getContextSnapshot ().getParentEndpoint ());
124+ ContextManager .stopSpan ();
125+ }
126+
127+ @ Test
128+ public void testOnResponse () throws InterruptedException {
129+ ContextManager .createEntrySpan ("sofarpc" , null );
130+ InvokeCallbackWrapper wrapper = new InvokeCallbackWrapper (callback );
131+ final WrapperWrapper wrapperWrapper = new WrapperWrapper (wrapper );
132+ executor .execute (() -> wrapperWrapper .onResponse (null ));
133+ ContextManager .stopSpan ();
134+ wrapperWrapper .getCountDownLatch ().await ();
135+
136+ assertThat (segmentStorage .getTraceSegments ().size (), is (2 ));
137+ TraceSegment traceSegment = segmentStorage .getTraceSegments ().get (0 );
138+ List <AbstractTracingSpan > spans = SegmentHelper .getSpans (traceSegment );
139+ assertThat (spans .size (), is (1 ));
140+
141+ TraceSegment traceSegment2 = segmentStorage .getTraceSegments ().get (1 );
142+ List <AbstractTracingSpan > spans2 = SegmentHelper .getSpans (traceSegment2 );
143+ assertThat (spans2 .size (), is (1 ));
144+ assertEquals ("sofarpc" , traceSegment2 .getRef ().getParentEndpoint ());
145+ }
146+
147+ @ Test
148+ public void testOnException () throws InterruptedException {
149+ ContextManager .createEntrySpan ("sofarpc" , null );
150+ InvokeCallbackWrapper wrapper = new InvokeCallbackWrapper (callback );
151+ final WrapperWrapper wrapperWrapper = new WrapperWrapper (wrapper );
152+ final Throwable throwable = new Throwable ();
153+ executor .execute (() -> wrapperWrapper .onException (throwable ));
154+ ContextManager .stopSpan ();
155+ wrapperWrapper .getCountDownLatch ().await ();
156+
157+ assertThat (segmentStorage .getTraceSegments ().size (), is (2 ));
158+ TraceSegment traceSegment = segmentStorage .getTraceSegments ().get (0 );
159+ List <AbstractTracingSpan > spans = SegmentHelper .getSpans (traceSegment );
160+ assertThat (spans .size (), is (1 ));
161+
162+ TraceSegment traceSegment2 = segmentStorage .getTraceSegments ().get (1 );
163+ List <AbstractTracingSpan > spans2 = SegmentHelper .getSpans (traceSegment2 );
164+ assertThat (spans2 .size (), is (1 ));
165+ assertThat (SpanHelper .getLogs (spans2 .get (0 )).size (), is (1 ));
166+
167+ }
168+
169+ }
0 commit comments