3939import static org .junit .Assert .assertEquals ;
4040import static org .junit .Assert .assertFalse ;
4141import static org .junit .Assert .assertNotNull ;
42+ import static org .junit .Assert .assertNull ;
4243import static org .junit .Assert .assertTrue ;
4344import static org .junit .Assert .fail ;
4445
@@ -179,7 +180,7 @@ address, NioSocketChannel.class, channelOptions, group, newNegotiator(),
179180 public void overrideDefaultUserAgent () throws Exception {
180181 startServer ();
181182 NettyClientTransport transport = newTransport (newNegotiator (),
182- DEFAULT_MAX_MESSAGE_SIZE , GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , "testUserAgent" );
183+ DEFAULT_MAX_MESSAGE_SIZE , GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , "testUserAgent" , true );
183184 callMeMaybe (transport .start (clientTransportListener ));
184185
185186 new Rpc (transport , new Metadata ()).halfClose ().waitForResponse ();
@@ -196,7 +197,7 @@ public void maxMessageSizeShouldBeEnforced() throws Throwable {
196197 startServer ();
197198 // Allow the response payloads of up to 1 byte.
198199 NettyClientTransport transport = newTransport (newNegotiator (),
199- 1 , GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , null );
200+ 1 , GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , null , true );
200201 callMeMaybe (transport .start (clientTransportListener ));
201202
202203 try {
@@ -278,7 +279,7 @@ public void maxHeaderListSizeShouldBeEnforcedOnClient() throws Exception {
278279 startServer ();
279280
280281 NettyClientTransport transport =
281- newTransport (newNegotiator (), DEFAULT_MAX_MESSAGE_SIZE , 1 , null );
282+ newTransport (newNegotiator (), DEFAULT_MAX_MESSAGE_SIZE , 1 , null , true );
282283 callMeMaybe (transport .start (clientTransportListener ));
283284
284285 try {
@@ -344,6 +345,24 @@ public void clientStreamGetsAttributes() throws Exception {
344345 assertEquals (address , rpc .stream .getAttributes ().get (Grpc .TRANSPORT_ATTR_REMOTE_ADDR ));
345346 }
346347
348+ @ Test
349+ public void keepAliveEnabled () throws Exception {
350+ startServer ();
351+ NettyClientTransport transport = newTransport (newNegotiator (), DEFAULT_MAX_MESSAGE_SIZE ,
352+ GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , null /* user agent */ , true /* keep alive */ );
353+ callMeMaybe (transport .start (clientTransportListener ));
354+ assertNotNull (transport .keepAliveManager ());
355+ }
356+
357+ @ Test
358+ public void keepAliveDisabled () throws Exception {
359+ startServer ();
360+ NettyClientTransport transport = newTransport (newNegotiator (), DEFAULT_MAX_MESSAGE_SIZE ,
361+ GrpcUtil .DEFAULT_MAX_HEADER_LIST_SIZE , null /* user agent */ , false /* keep alive */ );
362+ callMeMaybe (transport .start (clientTransportListener ));
363+ assertNull (transport .keepAliveManager ());
364+ }
365+
347366 private Throwable getRootCause (Throwable t ) {
348367 if (t .getCause () == null ) {
349368 return t ;
@@ -359,15 +378,18 @@ private ProtocolNegotiator newNegotiator() throws IOException {
359378 }
360379
361380 private NettyClientTransport newTransport (ProtocolNegotiator negotiator ) {
362- return newTransport (negotiator ,
363- DEFAULT_MAX_MESSAGE_SIZE , GrpcUtil . DEFAULT_MAX_HEADER_LIST_SIZE , null /* user agent */ );
381+ return newTransport (negotiator , DEFAULT_MAX_MESSAGE_SIZE , GrpcUtil . DEFAULT_MAX_HEADER_LIST_SIZE ,
382+ null /* user agent */ , true /* keep alive */ );
364383 }
365384
366- private NettyClientTransport newTransport (
367- ProtocolNegotiator negotiator , int maxMsgSize , int maxHeaderListSize , String userAgent ) {
385+ private NettyClientTransport newTransport (ProtocolNegotiator negotiator , int maxMsgSize ,
386+ int maxHeaderListSize , String userAgent , boolean enableKeepAlive ) {
368387 NettyClientTransport transport = new NettyClientTransport (
369388 address , NioSocketChannel .class , new HashMap <ChannelOption <?>, Object >(), group , negotiator ,
370389 DEFAULT_WINDOW_SIZE , maxMsgSize , maxHeaderListSize , authority , userAgent );
390+ if (enableKeepAlive ) {
391+ transport .enableKeepAlive (true , 1000 , 1000 );
392+ }
371393 transports .add (transport );
372394 return transport ;
373395 }
0 commit comments