1818
1919import static org .junit .Assert .assertEquals ;
2020import static org .junit .Assert .assertNotEquals ;
21+ import static org .junit .Assert .assertNotNull ;
2122import static org .junit .Assert .assertTrue ;
2223
24+ import io .grpc .Metadata ;
25+ import io .grpc .MethodDescriptor .MethodType ;
26+ import io .grpc .ServerCall ;
27+ import io .grpc .ServerCallHandler ;
28+ import io .grpc .ServerInterceptor ;
29+ import io .grpc .Status ;
30+ import java .util .ArrayList ;
31+ import java .util .List ;
32+ import java .util .Map ;
33+ import java .util .Objects ;
34+ import java .util .Set ;
35+ import java .util .concurrent .ConcurrentHashMap ;
36+ import java .util .concurrent .CopyOnWriteArrayList ;
2337import java .util .regex .Matcher ;
24- import java .util .regex .Pattern ;
2538import org .junit .Test ;
2639import org .junit .runner .RunWith ;
2740import org .junit .runners .JUnit4 ;
2841
2942@ RunWith (JUnit4 .class )
3043public class XGoogSpannerRequestIdTest {
31- private static final Pattern REGEX_RAND_PROCESS_ID =
32- Pattern .compile ("1.([0-9a-z]{16})(\\ .\\ d+){3}\\ .(\\ d+)$" );
3344
3445 @ Test
3546 public void testEquals () {
@@ -48,7 +59,139 @@ public void testEquals() {
4859 @ Test
4960 public void testEnsureHexadecimalFormatForRandProcessID () {
5061 String str = XGoogSpannerRequestId .of (1 , 2 , 3 , 4 ).toString ();
51- Matcher m = XGoogSpannerRequestIdTest . REGEX_RAND_PROCESS_ID .matcher (str );
62+ Matcher m = XGoogSpannerRequestId . REGEX .matcher (str );
5263 assertTrue (m .matches ());
5364 }
65+
66+ public static class ServerHeaderEnforcer implements ServerInterceptor {
67+ private Map <String , CopyOnWriteArrayList <XGoogSpannerRequestId >> unaryResults ;
68+ private Map <String , CopyOnWriteArrayList <XGoogSpannerRequestId >> streamingResults ;
69+ private List <String > gotValues ;
70+ private Set <String > checkMethods ;
71+
72+ ServerHeaderEnforcer (Set <String > checkMethods ) {
73+ this .gotValues = new CopyOnWriteArrayList <String >();
74+ this .unaryResults =
75+ new ConcurrentHashMap <String , CopyOnWriteArrayList <XGoogSpannerRequestId >>();
76+ this .streamingResults =
77+ new ConcurrentHashMap <String , CopyOnWriteArrayList <XGoogSpannerRequestId >>();
78+ this .checkMethods = checkMethods ;
79+ }
80+
81+ @ Override
82+ public <ReqT , RespT > ServerCall .Listener <ReqT > interceptCall (
83+ ServerCall <ReqT , RespT > call ,
84+ final Metadata requestHeaders ,
85+ ServerCallHandler <ReqT , RespT > next ) {
86+ boolean isUnary = call .getMethodDescriptor ().getType () == MethodType .UNARY ;
87+ String methodName = call .getMethodDescriptor ().getFullMethodName ();
88+ String gotReqIdStr = requestHeaders .get (XGoogSpannerRequestId .REQUEST_HEADER_KEY );
89+ if (!this .checkMethods .contains (methodName )) {
90+ System .out .println (
91+ "\033 [35mBypassing " + methodName + " but has " + gotReqIdStr + "\033 [00m" );
92+ return next .startCall (call , requestHeaders );
93+ }
94+
95+ Map <String , CopyOnWriteArrayList <XGoogSpannerRequestId >> saver = this .streamingResults ;
96+ if (isUnary ) {
97+ saver = this .unaryResults ;
98+ }
99+
100+ if (Objects .equals (gotReqIdStr , null ) || Objects .equals (gotReqIdStr , "" )) {
101+ Status status =
102+ Status .fromCode (Status .Code .INVALID_ARGUMENT )
103+ .augmentDescription (
104+ methodName + " lacks " + XGoogSpannerRequestId .REQUEST_HEADER_KEY );
105+ call .close (status , requestHeaders );
106+ return next .startCall (call , requestHeaders );
107+ }
108+
109+ assertNotNull (gotReqIdStr );
110+ // Firstly assert and validate that at least we've got a requestId.
111+ Matcher m = XGoogSpannerRequestId .REGEX .matcher (gotReqIdStr );
112+ assertTrue (m .matches ());
113+
114+ XGoogSpannerRequestId reqId = XGoogSpannerRequestId .of (gotReqIdStr );
115+ if (!saver .containsKey (methodName )) {
116+ saver .put (methodName , new CopyOnWriteArrayList <XGoogSpannerRequestId >());
117+ }
118+
119+ saver .get (methodName ).add (reqId );
120+
121+ // Finally proceed with the call.
122+ return next .startCall (call , requestHeaders );
123+ }
124+
125+ public String [] accumulatedValues () {
126+ return this .gotValues .toArray (new String [0 ]);
127+ }
128+
129+ public void assertIntegrity () {
130+ this .unaryResults .forEach (
131+ (String method , CopyOnWriteArrayList <XGoogSpannerRequestId > values ) -> {
132+ // System.out.println("\033[36munary.method: " + method + "\033[00m");
133+ XGoogSpannerRequestId .assertMonotonicityOfIds (method , values );
134+ });
135+ this .streamingResults .forEach (
136+ (String method , CopyOnWriteArrayList <XGoogSpannerRequestId > values ) -> {
137+ // System.out.println("\033[36mstreaming.method: " + method + "\033[00m");
138+ XGoogSpannerRequestId .assertMonotonicityOfIds (method , values );
139+ });
140+ }
141+
142+ public static class methodAndRequestId {
143+ String method ;
144+ String requestId ;
145+
146+ public methodAndRequestId (String method , String requestId ) {
147+ this .method = method ;
148+ this .requestId = requestId ;
149+ }
150+
151+ public String toString () {
152+ return "{" + this .method + ":" + this .requestId + "}" ;
153+ }
154+ }
155+
156+ public methodAndRequestId [] accumulatedUnaryValues () {
157+ List <methodAndRequestId > accumulated = new ArrayList ();
158+ this .unaryResults .forEach (
159+ (String method , CopyOnWriteArrayList <XGoogSpannerRequestId > values ) -> {
160+ for (int i = 0 ; i < values .size (); i ++) {
161+ accumulated .add (new methodAndRequestId (method , values .get (i ).toString ()));
162+ }
163+ });
164+ return accumulated .toArray (new methodAndRequestId [0 ]);
165+ }
166+
167+ public methodAndRequestId [] accumulatedStreamingValues () {
168+ List <methodAndRequestId > accumulated = new ArrayList ();
169+ this .streamingResults .forEach (
170+ (String method , CopyOnWriteArrayList <XGoogSpannerRequestId > values ) -> {
171+ for (int i = 0 ; i < values .size (); i ++) {
172+ accumulated .add (new methodAndRequestId (method , values .get (i ).toString ()));
173+ }
174+ });
175+ return accumulated .toArray (new methodAndRequestId [0 ]);
176+ }
177+
178+ public void printAccumulatedValues () {
179+ methodAndRequestId [] unary = this .accumulatedUnaryValues ();
180+ System .out .println ("accumulatedUnaryvalues" );
181+ for (int i = 0 ; i < unary .length ; i ++) {
182+ System .out .println ("\t " + unary [i ].toString ());
183+ }
184+ methodAndRequestId [] streaming = this .accumulatedStreamingValues ();
185+ System .out .println ("accumulatedStreaminvalues" );
186+ for (int i = 0 ; i < streaming .length ; i ++) {
187+ System .out .println ("\t " + streaming [i ].toString ());
188+ }
189+ }
190+
191+ public void reset () {
192+ this .gotValues .clear ();
193+ this .unaryResults .clear ();
194+ this .streamingResults .clear ();
195+ }
196+ }
54197}
0 commit comments