summer-ji-eng · Dec 10, 2020 · Dec 2, 2020 · Nov 25, 2020 · Nov 25, 2020 · Nov 26, 2020
diff --git a/src/main/java/com/google/api/generator/gapic/composer/ServiceClientClassComposer.java b/src/main/java/com/google/api/generator/gapic/composer/ServiceClientClassComposer.java
@@ -565,15 +565,10 @@ private static List<MethodDefinition> createMethodVariants(
               .setReturnType(methodOutputType)
               .build();
 
-      Optional<String> methodSampleCode = Optional.empty();
-      if (!method.hasLro()) {
-        // TODO(summerji): Remove the condition check once finished the implementation on lro sample
-        // code.
-        methodSampleCode =
-            Optional.of(
-                ServiceClientSampleCodeComposer.composeRpcMethodHeaderSampleCode(
-                    method, types.get(clientName), signature, resourceNames, messageTypes));
-      }
+      Optional<String> methodSampleCode =
+          Optional.of(
+              ServiceClientSampleCodeComposer.composeRpcMethodHeaderSampleCode(
+                  method, types.get(clientName), signature, resourceNames, messageTypes));
       MethodDefinition.Builder methodVariantBuilder =
           MethodDefinition.builder()
               .setHeaderCommentStatements(

diff --git a/src/main/java/com/google/api/generator/gapic/composer/ServiceClientSampleCodeComposer.java b/src/main/java/com/google/api/generator/gapic/composer/ServiceClientSampleCodeComposer.java
@@ -210,6 +210,10 @@ public static String composeRpcMethodHeaderSampleCode(
           composeUnaryPagedRpcMethodSampleCode(
               method, clientType, repeatedResponseType, arguments, resourceNames));
     }
+    if (method.hasLro()) {
+      return SampleCodeWriter.write(
+          composeUnaryLroRpcMethodSampleCode(method, clientType, arguments, resourceNames));
+    }
     return SampleCodeWriter.write(
         composeUnaryRpcMethodSampleCode(method, clientType, arguments, resourceNames));
   }
@@ -334,6 +338,68 @@ static TryCatchStatement composeUnaryPagedRpcMethodSampleCode(
         .build();
   }
 
+  @VisibleForTesting
+  static TryCatchStatement composeUnaryLroRpcMethodSampleCode(
+      Method method,
+      TypeNode clientType,
+      List<MethodArgument> arguments,
+      Map<String, ResourceName> resourceNames) {
+    VariableExpr clientVarExpr =
+        VariableExpr.withVariable(
+            Variable.builder()
+                .setName(JavaStyle.toLowerCamelCase(clientType.reference().name()))
+                .setType(clientType)
+                .build());
+    List<VariableExpr> rpcMethodArgVarExprs = createRpcMethodArgumentVariableExprs(arguments);
+    List<Expr> rpcMethodArgDefaultValueExprs =
+        createRpcMethodArgumentDefaultValueExprs(arguments, resourceNames);
+    List<Expr> bodyExprs =
+        createAssignmentsForVarExprsWithValueExprs(
+            rpcMethodArgVarExprs, rpcMethodArgDefaultValueExprs);
+    // Assign response variable with invoking client's lro method.
+    // e.g. if return void, echoClient.waitAsync(ttl).get(); or,
+    // e.g. if return other type, WaitResponse response = echoClient.waitAsync(ttl).get();
+    Expr invokeLroMethodExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(clientVarExpr)
+            .setMethodName(String.format("%sAsync", JavaStyle.toLowerCamelCase(method.name())))
+            .setArguments(
+                rpcMethodArgVarExprs.stream().map(e -> (Expr) e).collect(Collectors.toList()))
+            .build();
+    Expr getResponseMethodExpr =
+        MethodInvocationExpr.builder()
+            .setExprReferenceExpr(invokeLroMethodExpr)
+            .setMethodName("get")
+            .setReturnType(method.lro().responseType())
+            .build();
+    boolean returnsVoid = isProtoEmptyType(method.lro().responseType());
+    if (returnsVoid) {
+      bodyExprs.add(getResponseMethodExpr);
+    } else {
+      VariableExpr responseVarExpr =
+          VariableExpr.builder()
+              .setVariable(
+                  Variable.builder()
+                      .setName("response")
+                      .setType(method.lro().responseType())
+                      .build())
+              .setIsDecl(true)
+              .build();
+      bodyExprs.add(
+          AssignmentExpr.builder()
+              .setVariableExpr(responseVarExpr)
+              .setValueExpr(getResponseMethodExpr)
+              .build());
+    }
+
+    return TryCatchStatement.builder()
+        .setTryResourceExpr(assignClientVariableWithCreateMethodExpr(clientVarExpr))
+        .setTryBody(
+            bodyExprs.stream().map(e -> ExprStatement.withExpr(e)).collect(Collectors.toList()))
+        .setIsSampleCode(true)
+        .build();
+  }
+
   // ==================================Helpers===================================================//
 
   // Create a list of RPC method arguments' variable expressions.

diff --git a/...est/java/com/google/api/generator/gapic/composer/ServiceClientSampleCodeComposerTest.java b/...est/java/com/google/api/generator/gapic/composer/ServiceClientSampleCodeComposerTest.java
@@ -23,6 +23,7 @@
 import com.google.api.generator.engine.ast.VaporReference;
 import com.google.api.generator.gapic.composer.samplecode.SampleCodeWriter;
 import com.google.api.generator.gapic.model.Field;
+import com.google.api.generator.gapic.model.LongrunningOperation;
 import com.google.api.generator.gapic.model.Message;
 import com.google.api.generator.gapic.model.Method;
 import com.google.api.generator.gapic.model.MethodArgument;
@@ -40,6 +41,8 @@
 
 public class ServiceClientSampleCodeComposerTest {
   private static final String SHOWCASE_PACKAGE_NAME = "com.google.showcase.v1beta1";
+  private static final String LRO_PACKAGE_NAME = "com.google.longrunning";
+  private static final String PROTO_PACKAGE_NAME = "com.google.protobuf";
 
   // =======================================RPC Method Header Sample Code=======================//
   @Test
@@ -128,6 +131,76 @@ public void validComposeRpcMethodHeaderSampleCode_pagedUnaryRpc() {
     assertEquals(expected, results);
   }
 
+  @Test
+  public void validComposeRpcMethodHeaderSampleCode_lroUnaryRpc() {
+    FileDescriptor echoFileDescriptor = EchoOuterClass.getDescriptor();
+    Map<String, ResourceName> resourceNames = Parser.parseResourceNames(echoFileDescriptor);
+    Map<String, Message> messageTypes = Parser.parseMessages(echoFileDescriptor);
+    TypeNode clientType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("EchoClient")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    TypeNode inputType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("WaitRequest")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    TypeNode outputType =
+        TypeNode.withReference(
+            VaporReference.builder().setName("Operation").setPakkage(LRO_PACKAGE_NAME).build());
+    TypeNode responseType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("WaitResponse")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    TypeNode metadataType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("WaitMetadata")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    LongrunningOperation lro = LongrunningOperation.withTypes(responseType, metadataType);
+    TypeNode ttlTypeNode =
+        TypeNode.withReference(
+            VaporReference.builder().setName("Duration").setPakkage(PROTO_PACKAGE_NAME).build());
+    MethodArgument ttl =
+        MethodArgument.builder()
+            .setName("ttl")
+            .setType(ttlTypeNode)
+            .setField(
+                Field.builder()
+                    .setName("ttl")
+                    .setType(ttlTypeNode)
+                    .setIsMessage(true)
+                    .setIsContainedInOneof(true)
+                    .build())
+            .build();
+    List<MethodArgument> arguments = Arrays.asList(ttl);
+    Method method =
+        Method.builder()
+            .setName("Wait")
+            .setInputType(inputType)
+            .setOutputType(outputType)
+            .setLro(lro)
+            .setMethodSignatures(Arrays.asList(arguments))
+            .build();
+
+    String results =
+        ServiceClientSampleCodeComposer.composeRpcMethodHeaderSampleCode(
+            method, clientType, arguments, resourceNames, messageTypes);
+    String expected =
+        LineFormatter.lines(
+            "try (EchoClient echoClient = EchoClient.create()) {\n",
+            "  Duration ttl = Duration.newBuilder().build();\n",
+            "  WaitResponse response = echoClient.waitAsync(ttl).get();\n",
+            "}");
+    assertEquals(results, expected);
+  }
+
   @Test
   public void invalidComposeRpcMethodHeaderSampleCode_noMatchedRepeatedResponseTypeInPagedMethod() {
     FileDescriptor echoFileDescriptor = EchoOuterClass.getDescriptor();
@@ -770,7 +843,7 @@ public void composeUnaryRpcMethodSampleCode_methodReturnVoid() {
                 .build());
     TypeNode outputType =
         TypeNode.withReference(
-            VaporReference.builder().setName("Empty").setPakkage("com.google.protobuf").build());
+            VaporReference.builder().setName("Empty").setPakkage(PROTO_PACKAGE_NAME).build());
     List<List<MethodArgument>> methodSignatures =
         Arrays.asList(
             Arrays.asList(
@@ -921,4 +994,142 @@ public void composeUnaryPagedRpcMethodSampleCode_noMethodArguments() {
             "}");
     assertEquals(results, expected);
   }
+
+  // ===================================Unary LRO RPC Method Sample Code ======================//
+  @Test
+  public void composeUnaryLroRpcMethodSampleCode_lroReturnResponseType() {
+    FileDescriptor echoFileDescriptor = EchoOuterClass.getDescriptor();
+    Map<String, ResourceName> resourceNames = Parser.parseResourceNames(echoFileDescriptor);
+    TypeNode clientType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("EchoClient")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    TypeNode inputType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("WaitRequest")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    TypeNode outputType =
+        TypeNode.withReference(
+            VaporReference.builder().setName("Operation").setPakkage(LRO_PACKAGE_NAME).build());
+    TypeNode responseType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("WaitResponse")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    TypeNode metadataType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("WaitMetadata")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    LongrunningOperation lro = LongrunningOperation.withTypes(responseType, metadataType);
+    TypeNode ttlTypeNode =
+        TypeNode.withReference(
+            VaporReference.builder().setName("Duration").setPakkage(PROTO_PACKAGE_NAME).build());
+    MethodArgument ttl =
+        MethodArgument.builder()
+            .setName("ttl")
+            .setType(ttlTypeNode)
+            .setField(
+                Field.builder()
+                    .setName("ttl")
+                    .setType(ttlTypeNode)
+                    .setIsMessage(true)
+                    .setIsContainedInOneof(true)
+                    .build())
+            .build();
+    List<MethodArgument> arguments = Arrays.asList(ttl);
+    Method method =
+        Method.builder()
+            .setName("Wait")
+            .setInputType(inputType)
+            .setOutputType(outputType)
+            .setLro(lro)
+            .setMethodSignatures(Arrays.asList(arguments))
+            .build();
+
+    String results =
+        SampleCodeWriter.write(
+            ServiceClientSampleCodeComposer.composeUnaryLroRpcMethodSampleCode(
+                method, clientType, arguments, resourceNames));
+    String expected =
+        LineFormatter.lines(
+            "try (EchoClient echoClient = EchoClient.create()) {\n",
+            "  Duration ttl = Duration.newBuilder().build();\n",
+            "  WaitResponse response = echoClient.waitAsync(ttl).get();\n",
+            "}");
+    assertEquals(results, expected);
+  }
+
+  @Test
+  public void composeUnaryLroRpcMethodSampleCode_lroReturnVoid() {
+    FileDescriptor echoFileDescriptor = EchoOuterClass.getDescriptor();
+    Map<String, ResourceName> resourceNames = Parser.parseResourceNames(echoFileDescriptor);
+    TypeNode clientType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("EchoClient")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    TypeNode inputType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("WaitRequest")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    TypeNode outputType =
+        TypeNode.withReference(
+            VaporReference.builder().setName("Operation").setPakkage(LRO_PACKAGE_NAME).build());
+    TypeNode responseType =
+        TypeNode.withReference(
+            VaporReference.builder().setName("Empty").setPakkage(PROTO_PACKAGE_NAME).build());
+    TypeNode metadataType =
+        TypeNode.withReference(
+            VaporReference.builder()
+                .setName("WaitMetadata")
+                .setPakkage(SHOWCASE_PACKAGE_NAME)
+                .build());
+    LongrunningOperation lro = LongrunningOperation.withTypes(responseType, metadataType);
+    TypeNode ttlTypeNode =
+        TypeNode.withReference(
+            VaporReference.builder().setName("Duration").setPakkage(PROTO_PACKAGE_NAME).build());
+    MethodArgument ttl =
+        MethodArgument.builder()
+            .setName("ttl")
+            .setType(ttlTypeNode)
+            .setField(
+                Field.builder()
+                    .setName("ttl")
+                    .setType(ttlTypeNode)
+                    .setIsMessage(true)
+                    .setIsContainedInOneof(true)
+                    .build())
+            .build();
+    List<MethodArgument> arguments = Arrays.asList(ttl);
+    Method method =
+        Method.builder()
+            .setName("Wait")
+            .setInputType(inputType)
+            .setOutputType(outputType)
+            .setLro(lro)
+            .setMethodSignatures(Arrays.asList(arguments))
+            .build();
+
+    String results =
+        SampleCodeWriter.write(
+            ServiceClientSampleCodeComposer.composeUnaryLroRpcMethodSampleCode(
+                method, clientType, arguments, resourceNames));
+    String expected =
+        LineFormatter.lines(
+            "try (EchoClient echoClient = EchoClient.create()) {\n",
+            "  Duration ttl = Duration.newBuilder().build();\n",
+            "  echoClient.waitAsync(ttl).get();\n",
+            "}");
+    assertEquals(results, expected);
+  }
 }
diff --git a/src/test/java/com/google/api/generator/gapic/composer/goldens/EchoClient.golden b/src/test/java/com/google/api/generator/gapic/composer/goldens/EchoClient.golden
@@ -403,6 +403,15 @@ public class EchoClient implements BackgroundResource {
 
   // AUTO-GENERATED DOCUMENTATION AND METHOD.
   /**
+   * Sample code:
+   *
+   * <pre>{@code
+   * try (EchoClient echoClient = EchoClient.create()) {
+   *   Duration ttl = Duration.newBuilder().build();
+   *   WaitResponse response = echoClient.waitAsync(ttl).get();
+   * }
+   * }</pre>
+   *
    * @param ttl
    * @throws com.google.api.gax.rpc.ApiException if the remote call fails
    */
@@ -413,6 +422,15 @@ public class EchoClient implements BackgroundResource {
 
   // AUTO-GENERATED DOCUMENTATION AND METHOD.
   /**
+   * Sample code:
+   *
+   * <pre>{@code
+   * try (EchoClient echoClient = EchoClient.create()) {
+   *   Timestamp endTime = Timestamp.newBuilder().build();
+   *   WaitResponse response = echoClient.waitAsync(endTime).get();
+   * }
+   * }</pre>
+   *
    * @param endTime
    * @throws com.google.api.gax.rpc.ApiException if the remote call fails
    */