Fix possible NPE in ProxyMethodExample (#2867)

alzimmermsft · web-flow · commit e66eb9134bb6 · 2024-07-19T10:17:33.000+08:00
diff --git a/javagen/src/main/java/com/azure/autorest/model/clientmodel/ProxyMethodExample.java b/javagen/src/main/java/com/azure/autorest/model/clientmodel/ProxyMethodExample.java
@@ -10,7 +10,6 @@
 import com.azure.core.util.CoreUtils;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import com.fasterxml.jackson.databind.SerializationFeature;
 import org.slf4j.Logger;
 
 import java.net.MalformedURLException;
@@ -23,21 +22,18 @@
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
-import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 public class ProxyMethodExample {
 
     private final Logger logger = new PluginLogger(Javagen.getPluginInstance(), ProxyMethodExample.class);
 
-    private static final ObjectMapper PRETTY_PRINTER = new ObjectMapper().enable(SerializationFeature.INDENT_OUTPUT);
     private static final ObjectMapper NORMAL_PRINTER = new ObjectMapper();
     private static final String SLASH = "/";
-    private static final String QUOTED_SLASH = Pattern.quote(SLASH);
 
     private static String tspDirectory = null;
+
     public static void setTspDirectory(String tspDirectory) {
         ProxyMethodExample.tspDirectory = tspDirectory;
     }
@@ -61,7 +57,7 @@ public Object getObjectValue() {
 
         /**
          * Gets the un-escaped query value.
-         *
+         * <p>
          * This is done by heuristic, and not guaranteed to be correct.
          *
          * @return the un-escaped query value
@@ -77,13 +73,9 @@ public Object getUnescapedQueryValue() {
         @Override
         public String toString() {
             try {
-                return "ParameterValue{" +
-                        "objectValue=" + PRETTY_PRINTER.writeValueAsString(objectValue) +
-                        '}';
+                return "ParameterValue{" + "objectValue=" + NORMAL_PRINTER.writeValueAsString(objectValue) + '}';
             } catch (JsonProcessingException e) {
-                return "ParameterValue{" +
-                        "objectValue=" + objectValue +
-                        '}';
+                return "ParameterValue{" + "objectValue=" + objectValue + '}';
             }
         }
 
@@ -110,9 +102,8 @@ public Response(int statusCode, Object response) {
                 Map<String, Object> responseMap = (Map<String, Object>) response;
                 if (responseMap.containsKey("headers") && responseMap.get("headers") instanceof Map) {
                     Map<String, Object> headersMap = (Map<String, Object>) responseMap.get("headers");
-                    headersMap.forEach((header, value) -> {
-                        httpHeaders.add(HttpHeaderName.fromString(header), value.toString());
-                    });
+                    headersMap.forEach(
+                        (header, value) -> httpHeaders.add(HttpHeaderName.fromString(header), String.valueOf(value)));
                 }
                 this.body = responseMap.getOrDefault("body", null);
             } else {
@@ -167,17 +158,11 @@ public String getJson(Object obj) {
         @Override
         public String toString() {
             try {
-                return "Response{" +
-                        "statusCode=" + statusCode +
-                        ", httpHeaders=" + httpHeaders +
-                        ", body=" + PRETTY_PRINTER.writeValueAsString(body) +
-                        '}';
+                return "Response{" + "statusCode=" + statusCode + ", httpHeaders=" + httpHeaders + ", body="
+                    + NORMAL_PRINTER.writeValueAsString(body) + '}';
             } catch (JsonProcessingException e) {
-                return "Response{" +
-                        "statusCode=" + statusCode +
-                        ", httpHeaders=" + httpHeaders +
-                        ", body=" + body +
-                        '}';
+                return "Response{" + "statusCode=" + statusCode + ", httpHeaders=" + httpHeaders + ", body=" + body
+                    + '}';
             }
         }
     }
@@ -206,18 +191,23 @@ public Map<Integer, Response> getResponses() {
     /**
      * @return the primary response
      */
-    public Optional<Response> getPrimaryResponse() {
+    public Response getPrimaryResponse() {
         if (responses.isEmpty()) {
-            return Optional.empty();
+            return null;
         }
 
-        Optional<Response> response = responses.values().stream()
-                .filter(r -> r.statusCode / 100 == 2)
-                .findFirst();
-        if (!response.isPresent()) {
-            response = responses.values().stream().findFirst();
+        Response firstResponse = null;
+        for (Response response : responses.values()) {
+            if (firstResponse == null) {
+                firstResponse = response;
+            }
+
+            if (response.statusCode / 100 == 2) {
+                return response;
+            }
         }
-        return response;
+
+        return firstResponse;
     }
 
     /**
@@ -229,8 +219,9 @@ public String getOriginalFile() {
 
     /**
      * Heuristically find relative path of the original file to the repository.
-     *
-     * For instance, "specification/resources/resource-manager/Microsoft.Authorization/stable/2020-09-01/examples/getDataPolicyManifest.json"
+     * <p>
+     * For instance,
+     * "specification/resources/resource-manager/Microsoft.Authorization/stable/2020-09-01/examples/getDataPolicyManifest.json"
      *
      * @return the relative path of the original file
      */
@@ -241,32 +232,29 @@ public String getRelativeOriginalFileName() {
                 URL url = new URI(originalFileName).toURL();
                 switch (url.getProtocol()) {
                     case "http":
-                    case "https":
-                    {
-                        String[] segments = url.getPath().split(QUOTED_SLASH);
+                    case "https": {
+                        String[] segments = url.getPath().split(SLASH);
                         if (segments.length > 3) {
                             // first 3 should be owner, name, branch
                             originalFileName = Arrays.stream(segments)
-                                    .filter(s -> !s.isEmpty())
-                                    .skip(3)
-                                    .collect(Collectors.joining(SLASH));
+                                .filter(s -> !s.isEmpty())
+                                .skip(3)
+                                .collect(Collectors.joining(SLASH));
                         }
                         break;
                     }
 
-                    case "file":
-                    {
+                    case "file": {
                         String relativeFileName = tspDirectory != null
-                                ? getRelativeOriginalFileNameForTsp(url)
-                                : getRelativeOriginalFileNameForSwagger(url);
+                            ? getRelativeOriginalFileNameForTsp(url)
+                            : getRelativeOriginalFileNameForSwagger(url);
                         if (relativeFileName != null) {
                             originalFileName = relativeFileName;
                         }
                         break;
                     }
 
-                    default:
-                    {
+                    default: {
                         logger.error("Unknown protocol in x-ms-original-file: '{}'", originalFileName);
                         break;
                     }
@@ -281,8 +269,10 @@ public String getRelativeOriginalFileName() {
 
     /**
      * identifier of the codesnippet label from codesnippet-maven-plugin
-     * @see <a href="https://github.com/Azure/azure-sdk-tools/blob/main/packages/java-packages/codesnippet-maven-plugin/README.md">codesnippet-maven-plugin</a>
+     *
      * @return the identifier of the codesnippet label that wraps around the example code
+     * @see <a
+     * href="https://github.com/Azure/azure-sdk-tools/blob/main/packages/java-packages/codesnippet-maven-plugin/README.md">codesnippet-maven-plugin</a>
      */
     public String getCodeSnippetIdentifier() {
         return codeSnippetIdentifier;
@@ -308,21 +298,19 @@ static String getRelativeOriginalFileNameForTsp(URL url) {
          * "specification/standbypool/StandbyPool.Management/examples/2023-12-01-preview/StandbyVirtualMachinePools_Update.json"
          */
         String originalFileName = null;
-        String[] directorySegments = tspDirectory.split(QUOTED_SLASH);
+        String[] directorySegments = tspDirectory.split(SLASH);
         String directoryLastSegment = directorySegments[directorySegments.length - 1];
         int sharedDirectorySegment = -1;
-        String[] segments = url.getPath().split(QUOTED_SLASH);
+        String[] segments = url.getPath().split(SLASH);
         for (int i = segments.length - 1; i >= 0; --i) {
             if (Objects.equals(directoryLastSegment, segments[i])) {
                 sharedDirectorySegment = i;
                 break;
             }
         }
         if (sharedDirectorySegment >= 0) {
-            originalFileName = Stream.concat(
-                    Arrays.stream(directorySegments),
-                    Arrays.stream(segments).skip(sharedDirectorySegment + 1)
-            ).collect(Collectors.joining(SLASH));
+            originalFileName = Stream.concat(Arrays.stream(directorySegments),
+                Arrays.stream(segments).skip(sharedDirectorySegment + 1)).collect(Collectors.joining(SLASH));
         }
         return originalFileName;
     }
@@ -334,7 +322,7 @@ static String getRelativeOriginalFileNameForSwagger(URL url) {
          * or "specification/<service>/data-plane"
          */
         String originalFileName = null;
-        String[] segments = url.getPath().split(QUOTED_SLASH);
+        String[] segments = url.getPath().split(SLASH);
         int resourceManagerOrDataPlaneSegmentIndex = -1;
         for (int i = 0; i < segments.length; ++i) {
             if ("resource-manager".equals(segments[i]) || "data-plane".equals(segments[i])) {
@@ -344,18 +332,15 @@ static String getRelativeOriginalFileNameForSwagger(URL url) {
         }
         if (resourceManagerOrDataPlaneSegmentIndex > 2) {
             originalFileName = Arrays.stream(segments)
-                    .skip(resourceManagerOrDataPlaneSegmentIndex - 2)
-                    .collect(Collectors.joining(SLASH));
+                .skip(resourceManagerOrDataPlaneSegmentIndex - 2)
+                .collect(Collectors.joining(SLASH));
         }
         return originalFileName;
     }
 
     @Override
     public String toString() {
-        return "ProxyMethodExample{" +
-                "parameters=" + parameters +
-                ", responses=" + responses +
-                '}';
+        return "ProxyMethodExample{" + "parameters=" + parameters + ", responses=" + responses + '}';
     }
 
     public static final class Builder {
diff --git a/javagen/src/main/java/com/azure/autorest/template/example/ClientMethodExampleWriter.java b/javagen/src/main/java/com/azure/autorest/template/example/ClientMethodExampleWriter.java
@@ -39,7 +39,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
@@ -96,9 +95,8 @@ public ClientMethodExampleWriter(ClientMethod method, String clientVarName, Prox
             methodBlock.line(methodInvocation.toString());
         };
         responseAssertionWriter = methodBlock -> {
-            Optional<ProxyMethodExample.Response> responseOpt = proxyMethodExample.getPrimaryResponse();
-            if (responseOpt.isPresent()) {
-                ProxyMethodExample.Response response = responseOpt.get();
+            ProxyMethodExample.Response response = proxyMethodExample.getPrimaryResponse();
+            if (response != null) {
                 IType returnType = method.getReturnValue().getType();
                 if (returnType instanceof GenericType) {
                     GenericType responseType = (GenericType) returnType;
diff --git a/javagen/src/main/java/com/azure/autorest/template/example/ProtocolExampleWriter.java b/javagen/src/main/java/com/azure/autorest/template/example/ProtocolExampleWriter.java
@@ -37,7 +37,6 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
@@ -239,9 +238,8 @@ public ProtocolExampleWriter(ProtocolExample protocolExample) {
         };
 
         this.assertionWriter = methodBlock -> {
-            Optional<ProxyMethodExample.Response> responseOpt = proxyMethodExample.getPrimaryResponse();
-            if (responseOpt.isPresent()) {
-                ProxyMethodExample.Response response = responseOpt.get();
+            ProxyMethodExample.Response response = proxyMethodExample.getPrimaryResponse();
+            if (response != null) {
                 IType returnType = method.getReturnValue().getType();
                 if (returnType instanceof GenericType) {
                     GenericType responseType = (GenericType) returnType;
