fix(mixins): handle unit tests for mixin pagination methods

src/main/java/com/google/api/generator/gapic/composer/ServiceClientTestClassComposer.java

@@ -140,11 +140,7 @@ public GapicClass generate(GapicContext context, Service service) {
     Map<String, Message> messageTypes = context.messages();
     String pakkage = service.pakkage();
     TypeStore typeStore = new TypeStore();
-    addDynamicTypes(service, typeStore);
-    for (Service mixinService : context.mixinServices()) {
-      addDynamicTypes(mixinService, typeStore);
-    }
-
+    addDynamicTypes(context, service, typeStore);
     String className = ClassNames.getServiceClientTestClassName(service);
     GapicClass.Kind kind = Kind.MAIN;
 
@@ -466,6 +462,7 @@ private static List<MethodDefinition> createTestMethods(
         javaMethods.add(
             createRpcTestMethod(
                 method,
+                service,
                 matchingService,
                 Collections.emptyList(),
                 0,
@@ -487,6 +484,7 @@ private static List<MethodDefinition> createTestMethods(
           javaMethods.add(
               createRpcTestMethod(
                   method,
+                  service,
                   matchingService,
                   method.methodSignatures().get(i),
                   i,
@@ -509,9 +507,27 @@ private static List<MethodDefinition> createTestMethods(
     return javaMethods;
   }
 
+  /**
+   * Creates a test method for a given RPC, e.g. createAssetTest.
+   *
+   * @param method the RPC for which this test method is created.
+   * @param apiService the host service under test.
+   * @param rpcService the service that {@code method} belongs to. This is not equal to {@code
+   *     apiService} only when {@code method} is a mixin, in which case {@code rpcService} is the
+   *     mixed-in service. If {@code apiService} and {@code rpcService} are different, they will be
+   *     used only for pagination. Otherwise, {@code rpcService} subsumes {@code apiService}.
+   * @param methodSignature the method signature of the RPC under test.
+   * @param variantIndex the nth variant of the RPC under test. This applies when we have
+   *     polymorphism due to the presence of several method signature annotations in the proto.
+   * @param isRequestArg whether the RPC variant under test take only the request proto message.
+   * @param classMemberVarExprs the class members in the generated test class.
+   * @param resourceNames the resource names available for use.
+   * @param messageTypes the proto message types available for use.
+   */
   private static MethodDefinition createRpcTestMethod(
       Method method,
-      Service service,
+      Service apiService,
+      Service rpcService,
       List<MethodArgument> methodSignature,
       int variantIndex,
       boolean isRequestArg,
@@ -520,14 +536,15 @@ private static MethodDefinition createRpcTestMethod(
       Map<String, Message> messageTypes) {
     if (!method.stream().equals(Method.Stream.NONE)) {
       return createStreamingRpcTestMethod(
-          service, method, classMemberVarExprs, resourceNames, messageTypes);
+          rpcService, method, classMemberVarExprs, resourceNames, messageTypes);
     }
     // Construct the expected response.
     TypeNode methodOutputType = method.hasLro() ? method.lro().responseType() : method.outputType();
     List<Expr> methodExprs = new ArrayList<>();
 
     TypeNode repeatedResponseType = null;
     VariableExpr responsesElementVarExpr = null;
+    String mockServiceVarName = getMockServiceVarName(rpcService);
     if (method.isPaged()) {
       Message methodOutputMessage = messageTypes.get(method.outputType().reference().simpleName());
       Field repeatedPagedResultsField = methodOutputMessage.findAndUnwrapFirstRepeatedField();
@@ -596,6 +613,7 @@ private static MethodDefinition createRpcTestMethod(
             .setVariableExpr(expectedResponseVarExpr.toBuilder().setIsDecl(true).build())
             .setValueExpr(expectedResponseValExpr)
             .build());
+
     if (method.hasLro()) {
       VariableExpr resultOperationVarExpr =
           VariableExpr.withVariable(
@@ -613,14 +631,14 @@ private static MethodDefinition createRpcTestMethod(
               .build());
       methodExprs.add(
           MethodInvocationExpr.builder()
-              .setExprReferenceExpr(classMemberVarExprs.get(getMockServiceVarName(service)))
+              .setExprReferenceExpr(classMemberVarExprs.get(mockServiceVarName))
               .setMethodName("addResponse")
               .setArguments(resultOperationVarExpr)
               .build());
     } else {
       methodExprs.add(
           MethodInvocationExpr.builder()
-              .setExprReferenceExpr(classMemberVarExprs.get(getMockServiceVarName(service)))
+              .setExprReferenceExpr(classMemberVarExprs.get(mockServiceVarName))
               .setMethodName("addResponse")
               .setArguments(expectedResponseVarExpr)
               .build());
@@ -675,7 +693,12 @@ private static MethodDefinition createRpcTestMethod(
         VariableExpr.withVariable(
             Variable.builder()
                 .setType(
-                    method.isPaged() ? getPagedResponseType(method, service) : methodOutputType)
+                    !method.isPaged()
+                        ? methodOutputType
+                        // If this method is a paginated mixin, use the host service, since
+                        // ServiceClient defines the paged response class and the mixed-in service
+                        // does not have a client.
+                        : getPagedResponseType(method, method.isMixin() ? apiService : rpcService))
                 .setName(method.isPaged() ? "pagedListResponse" : "actualResponse")
                 .build());
     Expr rpcJavaMethodInvocationExpr =
@@ -828,7 +851,7 @@ private static MethodDefinition createRpcTestMethod(
             .setVariableExpr(actualRequestsVarExpr.toBuilder().setIsDecl(true).build())
             .setValueExpr(
                 MethodInvocationExpr.builder()
-                    .setExprReferenceExpr(classMemberVarExprs.get(getMockServiceVarName(service)))
+                    .setExprReferenceExpr(classMemberVarExprs.get(mockServiceVarName))
                     .setMethodName("getRequests")
                     .setReturnType(actualRequestsVarExpr.type())
                     .build())
@@ -1021,6 +1044,8 @@ private static MethodDefinition createStreamingRpcTestMethod(
             .setVariableExpr(expectedResponseVarExpr.toBuilder().setIsDecl(true).build())
             .setValueExpr(expectedResponseValExpr)
             .build());
+
+    String mockServiceVarName = getMockServiceVarName(service);
     if (method.hasLro()) {
       VariableExpr resultOperationVarExpr =
           VariableExpr.withVariable(
@@ -1038,14 +1063,14 @@ private static MethodDefinition createStreamingRpcTestMethod(
               .build());
       methodExprs.add(
           MethodInvocationExpr.builder()
-              .setExprReferenceExpr(classMemberVarExprs.get(getMockServiceVarName(service)))
+              .setExprReferenceExpr(classMemberVarExprs.get(mockServiceVarName))
               .setMethodName("addResponse")
               .setArguments(resultOperationVarExpr)
               .build());
     } else {
       methodExprs.add(
           MethodInvocationExpr.builder()
-              .setExprReferenceExpr(classMemberVarExprs.get(getMockServiceVarName(service)))
+              .setExprReferenceExpr(classMemberVarExprs.get(mockServiceVarName))
               .setMethodName("addResponse")
               .setArguments(expectedResponseVarExpr)
               .build());
@@ -1251,7 +1276,19 @@ private static MethodDefinition createStreamingRpcTestMethod(
         .build();
   }
 
-  // TODO(imraleung): Reorder params.
+  /**
+   * Creates a test method to exercise exceptions for a given RPC, e.g. createAssetTest.
+   *
+   * @param method the RPC for which this test method is created.
+   * @param service the service that {@code method} belongs to.
+   * @param methodSignature the method signature of the RPC under test.
+   * @param variantIndex the nth variant of the RPC under test. This applies when we have
+   *     polymorphism due to the presence of several method signature annotations in the proto.
+   * @param classMemberVarExprs the class members in the generated test class.
+   * @param resourceNames the resource names available for use.
+   * @param messageTypes the proto message types available for use.
+   */
+  // TODO(miraleung): Reorder params.
   private static MethodDefinition createRpcExceptionTestMethod(
       Method method,
       Service service,
@@ -1759,7 +1796,7 @@ private static Map<String, TypeNode> createDefaultMethodNamesToTypes() {
     return javaMethodNameToReturnType;
   }
 
-  private static void addDynamicTypes(Service service, TypeStore typeStore) {
+  private static void addDynamicTypes(GapicContext context, Service service, TypeStore typeStore) {
     typeStore.putAll(
         service.pakkage(),
         Arrays.asList(
@@ -1775,6 +1812,19 @@ private static void addDynamicTypes(Service service, TypeStore typeStore) {
             .collect(Collectors.toList()),
         true,
         ClassNames.getServiceClientClassName(service));
+    for (Service mixinService : context.mixinServices()) {
+      typeStore.put(mixinService.pakkage(), ClassNames.getMockServiceClassName(mixinService));
+      for (Method mixinMethod : mixinService.methods()) {
+        if (!mixinMethod.isPaged()) {
+          continue;
+        }
+        typeStore.put(
+            service.pakkage(),
+            String.format(PAGED_RESPONSE_TYPE_NAME_PATTERN, mixinMethod.name()),
+            true,
+            ClassNames.getServiceClientClassName(service));
+      }
+    }
   }
 
   private static TypeNode getOperationCallSettingsType(Method protoMethod) {

test/integration/BUILD.bazel

@@ -12,7 +12,6 @@ load(
     "golden_update",
     "integration_test",
 )
-
 load("@rules_proto//proto:defs.bzl", "proto_library")
 
 package(default_visibility = ["//visibility:public"])
@@ -204,6 +203,7 @@ proto_library(
         "@com_google_googleapis//google/api:client_proto",
         "@com_google_googleapis//google/api:field_behavior_proto",
         "@com_google_googleapis//google/api:resource_proto",
+        "@com_google_googleapis//google/cloud/location:location_proto",
         "@com_google_googleapis//google/iam/v1:iam_policy_proto",
         "@com_google_googleapis//google/iam/v1:policy_proto",
         "@com_google_protobuf//:duration_proto",
@@ -219,6 +219,7 @@ proto_library_with_info(
     deps = [
         ":kms_proto",
         "@com_google_googleapis//google/cloud:common_resources_proto",
+        "@com_google_googleapis//google/cloud/location:location_proto",
         "@com_google_googleapis//google/iam/v1:iam_policy_proto",
         "@com_google_googleapis//google/iam/v1:policy_proto",
     ],
@@ -243,10 +244,12 @@ java_gapic_library(
     service_yaml = "apis/kms/v1/cloudkms_test_mixins_v1.yaml",
     test_deps = [
         ":kms_java_grpc",
+        "@com_google_googleapis//google/cloud/location:location_java_grpc",
         "@com_google_googleapis//google/iam/v1:iam_java_grpc",
     ],
     deps = [
         ":kms_java_proto",
+        "@com_google_googleapis//google/cloud/location:location_java_proto",
         "@com_google_googleapis//google/iam/v1:iam_java_proto",
     ],
 )
@@ -267,5 +270,7 @@ java_gapic_assembly_gradle_pkg(
         "@com_google_googleapis//google/cloud/kms/v1:kms_java_grpc",
         "@com_google_googleapis//google/cloud/kms/v1:kms_java_proto",
         "@com_google_googleapis//google/cloud/kms/v1:kms_proto",
+        "@com_google_googleapis//google/cloud/location:location_java_grpc",
+        "@com_google_googleapis//google/cloud/location:location_java_proto",
     ],
 )

test/integration/apis/kms/v1/cloudkms_test_mixins_v1.yaml

@@ -5,6 +5,7 @@ title: Cloud Key Management Service (KMS) API
 
 apis:
 - name: google.cloud.kms.v1.KeyManagementService
+- name: google.cloud.location.Locations
 - name: google.iam.v1.IAMPolicy
 
 types:

test/integration/goldens/kms/GrpcKeyManagementServiceStub.java

@@ -20,6 +20,7 @@
 import static com.google.cloud.kms.v1.KeyManagementServiceClient.ListCryptoKeysPagedResponse;
 import static com.google.cloud.kms.v1.KeyManagementServiceClient.ListImportJobsPagedResponse;
 import static com.google.cloud.kms.v1.KeyManagementServiceClient.ListKeyRingsPagedResponse;
+import static com.google.cloud.kms.v1.KeyManagementServiceClient.ListLocationsPagedResponse;
 
 import com.google.api.gax.core.BackgroundResource;
 import com.google.api.gax.core.BackgroundResourceAggregation;
@@ -64,6 +65,10 @@
 import com.google.cloud.kms.v1.UpdateCryptoKeyPrimaryVersionRequest;
 import com.google.cloud.kms.v1.UpdateCryptoKeyRequest;
 import com.google.cloud.kms.v1.UpdateCryptoKeyVersionRequest;
+import com.google.cloud.location.GetLocationRequest;
+import com.google.cloud.location.ListLocationsRequest;
+import com.google.cloud.location.ListLocationsResponse;
+import com.google.cloud.location.Location;
 import com.google.common.collect.ImmutableMap;
 import com.google.iam.v1.GetIamPolicyRequest;
 import com.google.iam.v1.Policy;
@@ -340,6 +345,25 @@ public class GrpcKeyManagementServiceStub extends KeyManagementServiceStub {
                   ProtoUtils.marshaller(TestIamPermissionsResponse.getDefaultInstance()))
               .build();
 
+  private static final MethodDescriptor<ListLocationsRequest, ListLocationsResponse>
+      listLocationsMethodDescriptor =
+          MethodDescriptor.<ListLocationsRequest, ListLocationsResponse>newBuilder()
+              .setType(MethodDescriptor.MethodType.UNARY)
+              .setFullMethodName("google.cloud.location.Locations/ListLocations")
+              .setRequestMarshaller(
+                  ProtoUtils.marshaller(ListLocationsRequest.getDefaultInstance()))
+              .setResponseMarshaller(
+                  ProtoUtils.marshaller(ListLocationsResponse.getDefaultInstance()))
+              .build();
+
+  private static final MethodDescriptor<GetLocationRequest, Location> getLocationMethodDescriptor =
+      MethodDescriptor.<GetLocationRequest, Location>newBuilder()
+          .setType(MethodDescriptor.MethodType.UNARY)
+          .setFullMethodName("google.cloud.location.Locations/GetLocation")
+          .setRequestMarshaller(ProtoUtils.marshaller(GetLocationRequest.getDefaultInstance()))
+          .setResponseMarshaller(ProtoUtils.marshaller(Location.getDefaultInstance()))
+          .build();
+
   private final UnaryCallable<ListKeyRingsRequest, ListKeyRingsResponse> listKeyRingsCallable;
   private final UnaryCallable<ListKeyRingsRequest, ListKeyRingsPagedResponse>
       listKeyRingsPagedCallable;
@@ -384,6 +408,10 @@ public class GrpcKeyManagementServiceStub extends KeyManagementServiceStub {
   private final UnaryCallable<SetIamPolicyRequest, Policy> setIamPolicyCallable;
   private final UnaryCallable<TestIamPermissionsRequest, TestIamPermissionsResponse>
       testIamPermissionsCallable;
+  private final UnaryCallable<ListLocationsRequest, ListLocationsResponse> listLocationsCallable;
+  private final UnaryCallable<ListLocationsRequest, ListLocationsPagedResponse>
+      listLocationsPagedCallable;
+  private final UnaryCallable<GetLocationRequest, Location> getLocationCallable;
 
   private final BackgroundResource backgroundResources;
   private final GrpcOperationsStub operationsStub;
@@ -784,6 +812,32 @@ public Map<String, String> extract(TestIamPermissionsRequest request) {
                       }
                     })
                 .build();
+    GrpcCallSettings<ListLocationsRequest, ListLocationsResponse> listLocationsTransportSettings =
+        GrpcCallSettings.<ListLocationsRequest, ListLocationsResponse>newBuilder()
+            .setMethodDescriptor(listLocationsMethodDescriptor)
+            .setParamsExtractor(
+                new RequestParamsExtractor<ListLocationsRequest>() {
+                  @Override
+                  public Map<String, String> extract(ListLocationsRequest request) {
+                    ImmutableMap.Builder<String, String> params = ImmutableMap.builder();
+                    params.put("name", String.valueOf(request.getName()));
+                    return params.build();
+                  }
+                })
+            .build();
+    GrpcCallSettings<GetLocationRequest, Location> getLocationTransportSettings =
+        GrpcCallSettings.<GetLocationRequest, Location>newBuilder()
+            .setMethodDescriptor(getLocationMethodDescriptor)
+            .setParamsExtractor(
+                new RequestParamsExtractor<GetLocationRequest>() {
+                  @Override
+                  public Map<String, String> extract(GetLocationRequest request) {
+                    ImmutableMap.Builder<String, String> params = ImmutableMap.builder();
+                    params.put("name", String.valueOf(request.getName()));
+                    return params.build();
+                  }
+                })
+            .build();
 
     this.listKeyRingsCallable =
         callableFactory.createUnaryCallable(
@@ -897,6 +951,15 @@ public Map<String, String> extract(TestIamPermissionsRequest request) {
             testIamPermissionsTransportSettings,
             settings.testIamPermissionsSettings(),
             clientContext);
+    this.listLocationsCallable =
+        callableFactory.createUnaryCallable(
+            listLocationsTransportSettings, settings.listLocationsSettings(), clientContext);
+    this.listLocationsPagedCallable =
+        callableFactory.createPagedCallable(
+            listLocationsTransportSettings, settings.listLocationsSettings(), clientContext);
+    this.getLocationCallable =
+        callableFactory.createUnaryCallable(
+            getLocationTransportSettings, settings.getLocationSettings(), clientContext);
 
     this.backgroundResources =
         new BackgroundResourceAggregation(clientContext.getBackgroundResources());
@@ -1068,6 +1131,22 @@ public UnaryCallable<SetIamPolicyRequest, Policy> setIamPolicyCallable() {
     return testIamPermissionsCallable;
   }
 
+  @Override
+  public UnaryCallable<ListLocationsRequest, ListLocationsResponse> listLocationsCallable() {
+    return listLocationsCallable;
+  }
+
+  @Override
+  public UnaryCallable<ListLocationsRequest, ListLocationsPagedResponse>
+      listLocationsPagedCallable() {
+    return listLocationsPagedCallable;
+  }
+
+  @Override
+  public UnaryCallable<GetLocationRequest, Location> getLocationCallable() {
+    return getLocationCallable;
+  }
+
   @Override
   public final void close() {
     shutdown();