1919package org .apache .skywalking .apm .plugin .rocketMQ .v4 ;
2020
2121import java .util .List ;
22+ import java .util .Map ;
23+
2224import org .apache .rocketmq .client .impl .CommunicationMode ;
2325import org .apache .rocketmq .common .message .Message ;
26+ import org .apache .rocketmq .common .message .MessageDecoder ;
2427import org .apache .rocketmq .common .protocol .header .SendMessageRequestHeader ;
28+ import org .apache .skywalking .apm .agent .core .context .SW8ExtensionCarrierItem ;
2529import org .apache .skywalking .apm .agent .test .tools .SegmentStorage ;
2630import org .junit .Before ;
2731import org .junit .Rule ;
4448
4549import static org .hamcrest .CoreMatchers .is ;
4650import static org .junit .Assert .assertThat ;
51+ import static org .junit .Assert .assertTrue ;
4752import static org .mockito .Matchers .anyString ;
53+ import static org .mockito .Mockito .RETURNS_DEEP_STUBS ;
54+ import static org .mockito .Mockito .doAnswer ;
55+ import static org .mockito .Mockito .mock ;
4856import static org .mockito .Mockito .verify ;
4957import static org .powermock .api .mockito .PowerMockito .when ;
5058
59+ import static org .apache .rocketmq .common .message .MessageDecoder .NAME_VALUE_SEPARATOR ;
60+ import static org .apache .rocketmq .common .message .MessageDecoder .PROPERTY_SEPARATOR ;
61+
5162@ RunWith (PowerMockRunner .class )
5263@ PowerMockRunnerDelegate (TracingSegmentRunner .class )
5364public class MessageSendInterceptorTest {
@@ -108,15 +119,22 @@ public void setSkyWalkingDynamicField(Object value) {
108119 CommunicationMode .ASYNC ,
109120 null
110121 };
111- when (messageRequestHeader .getProperties ()).thenReturn ("" );
112122 when (message .getTags ()).thenReturn ("TagA" );
123+ stubMessageRequestHeader ("TAGS" + NAME_VALUE_SEPARATOR + "TagA" + PROPERTY_SEPARATOR );
113124 }
114125
115126 @ Test
116127 public void testSendMessage () throws Throwable {
117128 messageSendInterceptor .beforeMethod (enhancedInstance , null , arguments , null , null );
118129 messageSendInterceptor .afterMethod (enhancedInstance , null , arguments , null , null );
119130
131+ Map <String , String > tags = MessageDecoder .string2messageProperties (
132+ ((SendMessageRequestHeader ) arguments [3 ]).getProperties ());
133+ // check original header of TAGS
134+ assertThat (tags .get ("TAGS" ), is ("TagA" ));
135+ // check skywalking header
136+ assertTrue (tags .containsKey (SW8ExtensionCarrierItem .HEADER_NAME ));
137+
120138 assertThat (segmentStorage .getTraceSegments ().size (), is (1 ));
121139 TraceSegment traceSegment = segmentStorage .getTraceSegments ().get (0 );
122140 List <AbstractTracingSpan > spans = SegmentHelper .getSpans (traceSegment );
@@ -127,15 +145,27 @@ public void testSendMessage() throws Throwable {
127145 SpanAssert .assertLayer (mqSpan , SpanLayer .MQ );
128146 SpanAssert .assertComponent (mqSpan , ComponentsDefine .ROCKET_MQ_PRODUCER );
129147 SpanAssert .assertTag (mqSpan , 0 , "127.0.0.1" );
130- verify (messageRequestHeader ).setProperties (anyString ());
131148 verify (callBack ).setSkyWalkingDynamicField (Matchers .any ());
132149 }
133150
151+ @ Test
152+ public void testSendMessageNew () throws Throwable {
153+ stubMessageRequestHeader ("TAGS" + NAME_VALUE_SEPARATOR + "TagA" );
154+ testSendMessage ();
155+ }
156+
134157 @ Test
135158 public void testSendMessageWithoutCallBack () throws Throwable {
136159 messageSendInterceptor .beforeMethod (enhancedInstance , null , argumentsWithoutCallback , null , null );
137160 messageSendInterceptor .afterMethod (enhancedInstance , null , argumentsWithoutCallback , null , null );
138161
162+ Map <String , String > tags = MessageDecoder .string2messageProperties (
163+ ((SendMessageRequestHeader ) argumentsWithoutCallback [3 ]).getProperties ());
164+ // check original header of TAGS
165+ assertThat (tags .get ("TAGS" ), is ("TagA" ));
166+ // check skywalking header
167+ assertTrue (tags .containsKey (SW8ExtensionCarrierItem .HEADER_NAME ));
168+
139169 assertThat (segmentStorage .getTraceSegments ().size (), is (1 ));
140170 TraceSegment traceSegment = segmentStorage .getTraceSegments ().get (0 );
141171 List <AbstractTracingSpan > spans = SegmentHelper .getSpans (traceSegment );
@@ -146,7 +176,25 @@ public void testSendMessageWithoutCallBack() throws Throwable {
146176 SpanAssert .assertLayer (mqSpan , SpanLayer .MQ );
147177 SpanAssert .assertComponent (mqSpan , ComponentsDefine .ROCKET_MQ_PRODUCER );
148178 SpanAssert .assertTag (mqSpan , 0 , "127.0.0.1" );
149- verify (messageRequestHeader ).setProperties (anyString ());
150179 }
151180
181+ @ Test
182+ public void testSendMessageWithoutCallBackNew () throws Throwable {
183+ stubMessageRequestHeader ("TAGS" + NAME_VALUE_SEPARATOR + "TagA" );
184+ testSendMessageWithoutCallBack ();
185+ }
186+
187+ private void stubMessageRequestHeader (String properties ) {
188+ messageRequestHeader = mock (SendMessageRequestHeader .class , RETURNS_DEEP_STUBS );
189+ doAnswer (invocation -> {
190+ String val = (String ) invocation .getArguments ()[0 ];
191+ when (messageRequestHeader .getProperties ()).thenReturn (val );
192+ return null ;
193+ }).when (messageRequestHeader ).setProperties (anyString ());
194+ when (messageRequestHeader .getProperties ()).thenCallRealMethod ();
195+ messageRequestHeader .setProperties (properties );
196+
197+ arguments [3 ] = messageRequestHeader ;
198+ argumentsWithoutCallback [3 ] = messageRequestHeader ;
199+ }
152200}
0 commit comments