@@ -492,27 +492,60 @@ private static MethodDefinition createRpcTestMethod(
492492 // Construct the expected response.
493493 // TODO(miraleung): Paging here.
494494 TypeNode methodOutputType = method .hasLro () ? method .lro ().responseType () : method .outputType ();
495+ List <Expr > methodExprs = new ArrayList <>();
496+
497+ TypeNode repeatedResponseType = null ;
498+ VariableExpr responsesElementVarExpr = null ;
499+ if (method .isPaged ()) {
500+ Message methodOutputMessage = messageTypes .get (method .outputType ().reference ().name ());
501+ repeatedResponseType = findRepeatedPagedType (methodOutputMessage );
502+ Preconditions .checkNotNull (
503+ repeatedResponseType ,
504+ String .format (
505+ "No repeated type found for paged method %s with output message type %s" ,
506+ method .name (), methodOutputMessage .name ()));
507+ responsesElementVarExpr =
508+ VariableExpr .withVariable (
509+ Variable .builder ().setType (repeatedResponseType ).setName ("responsesElement" ).build ());
510+ methodExprs .add (
511+ AssignmentExpr .builder ()
512+ .setVariableExpr (responsesElementVarExpr .toBuilder ().setIsDecl (true ).build ())
513+ .setValueExpr (
514+ DefaultValueComposer .createDefaultValue (
515+ Field .builder ()
516+ .setType (repeatedResponseType )
517+ .setName ("responsesElement" )
518+ .setIsMessage (true )
519+ .build ()))
520+ .build ());
521+ }
522+
495523 VariableExpr expectedResponseVarExpr =
496524 VariableExpr .withVariable (
497525 Variable .builder ().setType (methodOutputType ).setName ("expectedResponse" ).build ());
498526 Expr expectedResponseValExpr = null ;
499- if (messageTypes . containsKey ( methodOutputType . reference (). name () )) {
527+ if (method . isPaged ( )) {
500528 expectedResponseValExpr =
501- DefaultValueComposer .createSimpleMessageBuilderExpr (
502- messageTypes . get ( methodOutputType . reference (). name ()), resourceNames , messageTypes );
529+ DefaultValueComposer .createSimplePagedResponse (
530+ method . outputType (), responsesElementVarExpr );
503531 } else {
504- // Wrap this in a field so we don't have to split the helper into lots of different methods,
505- // or duplicate it for VariableExpr.
506- expectedResponseValExpr =
507- DefaultValueComposer .createDefaultValue (
508- Field .builder ()
509- .setType (methodOutputType )
510- .setIsMessage (true )
511- .setName ("expectedResponse" )
512- .build ());
532+ if (messageTypes .containsKey (methodOutputType .reference ().name ())) {
533+ expectedResponseValExpr =
534+ DefaultValueComposer .createSimpleMessageBuilderExpr (
535+ messageTypes .get (methodOutputType .reference ().name ()), resourceNames , messageTypes );
536+ } else {
537+ // Wrap this in a field so we don't have to split the helper into lots of different methods,
538+ // or duplicate it for VariableExpr.
539+ expectedResponseValExpr =
540+ DefaultValueComposer .createDefaultValue (
541+ Field .builder ()
542+ .setType (methodOutputType )
543+ .setIsMessage (true )
544+ .setName ("expectedResponse" )
545+ .build ());
546+ }
513547 }
514548
515- List <Expr > methodExprs = new ArrayList <>();
516549 methodExprs .add (
517550 AssignmentExpr .builder ()
518551 .setVariableExpr (expectedResponseVarExpr .toBuilder ().setIsDecl (true ).build ())
@@ -588,7 +621,10 @@ private static MethodDefinition createRpcTestMethod(
588621 // Call the RPC Java method.
589622 VariableExpr actualResponseVarExpr =
590623 VariableExpr .withVariable (
591- Variable .builder ().setType (methodOutputType ).setName ("actualResponse" ).build ());
624+ Variable .builder ()
625+ .setType (methodOutputType )
626+ .setName (method .isPaged () ? "pagedListResponse" : "actualResponse" )
627+ .build ());
592628 Expr rpcJavaMethodInvocationExpr =
593629 MethodInvocationExpr .builder ()
594630 .setExprReferenceExpr (classMemberVarExprs .get ("client" ))
@@ -610,12 +646,88 @@ private static MethodDefinition createRpcTestMethod(
610646 .setVariableExpr (actualResponseVarExpr .toBuilder ().setIsDecl (true ).build ())
611647 .setValueExpr (rpcJavaMethodInvocationExpr )
612648 .build ());
613- methodExprs .add (
614- MethodInvocationExpr .builder ()
615- .setStaticReferenceType (STATIC_TYPES .get ("Assert" ))
616- .setMethodName ("assertEquals" )
617- .setArguments (expectedResponseVarExpr , actualResponseVarExpr )
618- .build ());
649+
650+ if (method .isPaged ()) {
651+ // Assign the resources variaqble.
652+ VariableExpr resourcesVarExpr =
653+ VariableExpr .withVariable (
654+ Variable .builder ()
655+ .setType (
656+ TypeNode .withReference (
657+ ConcreteReference .builder ()
658+ .setClazz (List .class )
659+ .setGenerics (Arrays .asList (repeatedResponseType .reference ()))
660+ .build ()))
661+ .setName ("resources" )
662+ .build ());
663+ Expr iterateAllExpr =
664+ MethodInvocationExpr .builder ()
665+ .setExprReferenceExpr (actualResponseVarExpr )
666+ .setMethodName ("iterateAll" )
667+ .build ();
668+ Expr resourcesValExpr =
669+ MethodInvocationExpr .builder ()
670+ .setStaticReferenceType (STATIC_TYPES .get ("Lists" ))
671+ .setMethodName ("newArrayList" )
672+ .setArguments (iterateAllExpr )
673+ .setReturnType (resourcesVarExpr .type ())
674+ .build ();
675+ methodExprs .add (
676+ AssignmentExpr .builder ()
677+ .setVariableExpr (resourcesVarExpr )
678+ .setValueExpr (resourcesValExpr )
679+ .build ());
680+
681+ // Assert the size is equivalent.
682+ methodExprs .add (
683+ MethodInvocationExpr .builder ()
684+ .setStaticReferenceType (STATIC_TYPES .get ("Assert" ))
685+ .setMethodName ("assertEquals" )
686+ .setArguments (
687+ ValueExpr .withValue (
688+ PrimitiveValue .builder ().setType (TypeNode .INT ).setValue ("1" ).build ()),
689+ MethodInvocationExpr .builder ()
690+ .setExprReferenceExpr (resourcesVarExpr )
691+ .setMethodName ("size" )
692+ .setReturnType (TypeNode .INT )
693+ .build ())
694+ .build ());
695+
696+ // Assert the responses are equivalent.
697+ Expr zeroExpr =
698+ ValueExpr .withValue (PrimitiveValue .builder ().setType (TypeNode .INT ).setValue ("0" ).build ());
699+ Expr expectedPagedResponseExpr =
700+ MethodInvocationExpr .builder ()
701+ .setExprReferenceExpr (expectedResponseVarExpr )
702+ .setMethodName ("getResponsesList" )
703+ .build ();
704+ expectedPagedResponseExpr =
705+ MethodInvocationExpr .builder ()
706+ .setExprReferenceExpr (expectedPagedResponseExpr )
707+ .setMethodName ("get" )
708+ .setArguments (zeroExpr )
709+ .build ();
710+ Expr actualPagedResponseExpr =
711+ MethodInvocationExpr .builder ()
712+ .setExprReferenceExpr (resourcesVarExpr )
713+ .setMethodName ("get" )
714+ .setArguments (zeroExpr )
715+ .build ();
716+
717+ methodExprs .add (
718+ MethodInvocationExpr .builder ()
719+ .setStaticReferenceType (STATIC_TYPES .get ("Assert" ))
720+ .setMethodName ("assertEquals" )
721+ .setArguments (expectedPagedResponseExpr , actualPagedResponseExpr )
722+ .build ());
723+ } else {
724+ methodExprs .add (
725+ MethodInvocationExpr .builder ()
726+ .setStaticReferenceType (STATIC_TYPES .get ("Assert" ))
727+ .setMethodName ("assertEquals" )
728+ .setArguments (expectedResponseVarExpr , actualResponseVarExpr )
729+ .build ());
730+ }
619731 // TODO(miraleung): Empty line here.
620732
621733 // Construct the request checker logic.
@@ -658,7 +770,7 @@ private static MethodDefinition createRpcTestMethod(
658770
659771 VariableExpr actualRequestVarExpr =
660772 VariableExpr .withVariable (
661- Variable .builder ().setType (methodOutputType ).setName ("actualRequest" ).build ());
773+ Variable .builder ().setType (method . inputType () ).setName ("actualRequest" ).build ());
662774 Expr getFirstRequestExpr =
663775 MethodInvocationExpr .builder ()
664776 .setExprReferenceExpr (actualRequestsVarExpr )
@@ -669,7 +781,7 @@ private static MethodDefinition createRpcTestMethod(
669781 .setReturnType (STATIC_TYPES .get ("AbstractMessage" ))
670782 .build ();
671783 getFirstRequestExpr =
672- CastExpr .builder ().setType (methodOutputType ).setExpr (getFirstRequestExpr ).build ();
784+ CastExpr .builder ().setType (method . inputType () ).setExpr (getFirstRequestExpr ).build ();
673785 methodExprs .add (
674786 AssignmentExpr .builder ()
675787 .setVariableExpr (actualRequestVarExpr .toBuilder ().setIsDecl (true ).build ())
@@ -1437,6 +1549,16 @@ private static TypeNode getCallableType(Method protoMethod) {
14371549 ConcreteReference .builder ().setClazz (callableClazz ).setGenerics (generics ).build ());
14381550 }
14391551
1552+ private static TypeNode findRepeatedPagedType (Message message ) {
1553+ for (Field field : message .fields ()) {
1554+ if (field .isRepeated () && !field .isMap ()) {
1555+ Reference repeatedGenericRef = field .type ().reference ().generics ().get (0 );
1556+ return TypeNode .withReference (repeatedGenericRef );
1557+ }
1558+ }
1559+ return null ;
1560+ }
1561+
14401562 private static String getCallableMethodName (Method protoMethod ) {
14411563 Preconditions .checkState (
14421564 !protoMethod .stream ().equals (Method .Stream .NONE ),
0 commit comments