mgmt, support more type for generated mock test

fluentgen/src/main/java/com/azure/autorest/fluent/mapper/ExampleParser.java

@@ -162,7 +162,7 @@ private static List<FluentCollectionMethodExample> parseMethod(FluentResourceCol
         List<FluentCollectionMethodExample> ret = null;
 
         ClientMethod clientMethod = collectionMethod.getInnerClientMethod();
-        if (FluentUtils.validToGenerateExample(clientMethod)) {
+        if (FluentUtils.validRequestContentTypeToGenerateExample(clientMethod)) {
             ret = new ArrayList<>();
 
             List<MethodParameter> methodParameters = getParameters(clientMethod);
@@ -180,7 +180,7 @@ private static List<FluentCollectionMethodExample> parseMethod(FluentResourceCol
     private static List<FluentClientMethodExample> parseMethod(MethodGroupClient methodGroup, ClientMethod clientMethod) {
         List<FluentClientMethodExample> ret = null;
 
-        if (FluentUtils.validToGenerateExample(clientMethod)) {
+        if (FluentUtils.validRequestContentTypeToGenerateExample(clientMethod)) {
             ret = new ArrayList<>();
 
             List<MethodParameter> methodParameters = getParameters(clientMethod);
@@ -245,7 +245,7 @@ private static List<FluentResourceCreateExample> parseResourceCreate(FluentResou
         List<FluentCollectionMethod> collectionMethods = resourceCreate.getMethodReferences();
         for (FluentCollectionMethod collectionMethod : collectionMethods) {
             ClientMethod clientMethod = collectionMethod.getInnerClientMethod();
-            if (FluentUtils.validToGenerateExample(clientMethod)) {
+            if (FluentUtils.validRequestContentTypeToGenerateExample(clientMethod)) {
                 if (ret == null) {
                     ret = new ArrayList<>();
                 }
@@ -361,7 +361,7 @@ private static List<FluentResourceUpdateExample> parseResourceUpdate(FluentResou
         List<FluentCollectionMethod> collectionMethods = resourceUpdate.getMethodReferences();
         for (FluentCollectionMethod collectionMethod : collectionMethods) {
             ClientMethod clientMethod = collectionMethod.getInnerClientMethod();
-            if (FluentUtils.validToGenerateExample(clientMethod)) {
+            if (FluentUtils.validRequestContentTypeToGenerateExample(clientMethod)) {
                 if (ret == null) {
                     ret = new ArrayList<>();
                 }

fluentgen/src/main/java/com/azure/autorest/fluent/mapper/MockUnitTestParser.java

@@ -62,7 +62,7 @@ private static FluentMethodMockUnitTest parseResourceCreate(FluentResourceCollec
             List<FluentCollectionMethod> collectionMethods = resourceCreate.getMethodReferences();
             for (FluentCollectionMethod collectionMethod : collectionMethods) {
                 ClientMethod clientMethod = collectionMethod.getInnerClientMethod();
-                if (FluentUtils.validToGenerateExample(clientMethod) && requiresExample(clientMethod)) {
+                if (FluentUtils.validRequestContentTypeToGenerateExample(clientMethod) && FluentUtils.validResponseContentTypeToGenerateExample(clientMethod) && requiresExample(clientMethod)) {
                     List<MethodParameter> methodParameters = getParameters(clientMethod);
                     MethodParameter requestBodyParameter = findRequestBodyParameter(methodParameters);
                     ProxyMethodExample proxyMethodExample = createProxyMethodExample(clientMethod, methodParameters);
@@ -71,6 +71,7 @@ private static FluentMethodMockUnitTest parseResourceCreate(FluentResourceCollec
 
                     ResponseInfo responseInfo = createProxyMethodExampleResponse(clientMethod);
                     unitTest = new FluentMethodMockUnitTest(resourceCreateExample, collection, collectionMethod,
+                            FluentUtils.isResponseType(collectionMethod.getFluentReturnType()) ? FluentUtils.getValueTypeFromResponseType(collectionMethod.getFluentReturnType()) : collectionMethod.getFluentReturnType(),
                             responseInfo.responseExample, responseInfo.verificationObjectName, responseInfo.verificationNode);
 
                     break;
@@ -87,14 +88,14 @@ private static FluentMethodMockUnitTest parseMethod(FluentResourceCollection col
 
         try {
             ClientMethod clientMethod = collectionMethod.getInnerClientMethod();
-            if (FluentUtils.validToGenerateExample(clientMethod) && requiresExample(clientMethod)) {
+            if (FluentUtils.validRequestContentTypeToGenerateExample(clientMethod) && FluentUtils.validResponseContentTypeToGenerateExample(clientMethod) && requiresExample(clientMethod)) {
                 List<MethodParameter> methodParameters = getParameters(clientMethod);
                 ProxyMethodExample proxyMethodExample = createProxyMethodExample(clientMethod, methodParameters);
                 FluentCollectionMethodExample collectionMethodExample =
                         parseMethodForExample(collection, collectionMethod, methodParameters, proxyMethodExample.getName(), proxyMethodExample);
 
                 ResponseInfo responseInfo = createProxyMethodExampleResponse(clientMethod);
-                unitTest = new FluentMethodMockUnitTest(collectionMethodExample, collection, collectionMethod,
+                unitTest = new FluentMethodMockUnitTest(collectionMethodExample, collection, collectionMethod, collectionMethod.getFluentReturnType(),
                         responseInfo.responseExample, responseInfo.verificationObjectName, responseInfo.verificationNode);
             }
         } catch (PossibleCredentialException e) {
@@ -143,10 +144,17 @@ private static ResponseInfo createProxyMethodExampleResponse(ClientMethod client
         Object jsonObject;
         ExampleNode verificationNode;
         String verificationObjectName;
+
+        IType clientReturnType = clientMethod.getReturnValue().getType();
+        final boolean isResponseType = FluentUtils.isResponseType(clientReturnType);
+        if (isResponseType) {
+            clientReturnType = FluentUtils.getValueTypeFromResponseType(clientReturnType);
+        }
+
         if (clientMethod.getType() == ClientMethodType.PagingSync) {
             // pageable
-            if (clientMethod.getReturnValue().getType() instanceof GenericType) {
-                IType elementType = ((GenericType) clientMethod.getReturnValue().getType()).getTypeArguments()[0];
+            if (clientReturnType instanceof GenericType) {
+                IType elementType = ((GenericType) clientReturnType).getTypeArguments()[0];
 
                 Object firstJsonObjectInPageable = ModelTestCaseUtil.jsonFromType(0, elementType);
                 // put to first element in array
@@ -163,7 +171,7 @@ private static ResponseInfo createProxyMethodExampleResponse(ClientMethod client
             }
         } else {
             // simple or LRO
-            jsonObject = ModelTestCaseUtil.jsonFromType(0, clientMethod.getReturnValue().getType());
+            jsonObject = ModelTestCaseUtil.jsonFromType(0, clientReturnType);
 
             if (jsonObject == null) {
                 jsonObject = new Object();
@@ -174,7 +182,7 @@ private static ResponseInfo createProxyMethodExampleResponse(ClientMethod client
             }
 
             verificationObjectName = "response";
-            verificationNode = ModelExampleUtil.parseNode(clientMethod.getReturnValue().getType(), jsonObject);
+            verificationNode = ModelExampleUtil.parseNode(clientReturnType, jsonObject);
         }
         Map<String, Object> responseObject = new HashMap<>();
         responseObject.put("body", jsonObject);
@@ -183,6 +191,7 @@ private static ResponseInfo createProxyMethodExampleResponse(ClientMethod client
 
     private static boolean requiresExample(ClientMethod clientMethod) {
         if (clientMethod.getType() == ClientMethodType.SimpleSync
+                || clientMethod.getType() == ClientMethodType.SimpleSyncRestResponse
                 || clientMethod.getType() == ClientMethodType.PagingSync
                 // limit the scope of LRO to status code of 200
                 || (clientMethod.getType() == ClientMethodType.LongRunningSync

fluentgen/src/main/java/com/azure/autorest/fluent/model/clientmodel/examplemodel/FluentMethodMockUnitTest.java

@@ -5,13 +5,15 @@
 
 import com.azure.autorest.fluent.model.clientmodel.FluentCollectionMethod;
 import com.azure.autorest.fluent.model.clientmodel.FluentResourceCollection;
+import com.azure.autorest.model.clientmodel.IType;
 import com.azure.autorest.model.clientmodel.ProxyMethodExample;
 import com.azure.autorest.model.clientmodel.examplemodel.ExampleNode;
 
 public class FluentMethodMockUnitTest {
 
     // method with mock data
     private final FluentMethodExample fluentMethodExample;
+    private final IType fluentReturnType;
     private final FluentResourceCreateExample fluentResourceCreateExample;
 
     private final FluentResourceCollection resourceCollection;
@@ -26,13 +28,14 @@ public class FluentMethodMockUnitTest {
     public FluentMethodMockUnitTest(
             FluentMethodExample fluentMethodExample,
             FluentResourceCollection resourceCollection, FluentCollectionMethod fluentCollectionMethod,
-            ProxyMethodExample.Response response,
+            IType fluentReturnType, ProxyMethodExample.Response response,
             String responseVerificationVariableName, ExampleNode responseVerificationNode) {
 
         this.fluentMethodExample = fluentMethodExample;
         this.fluentResourceCreateExample = null;
         this.resourceCollection = resourceCollection;
         this.collectionMethod = fluentCollectionMethod;
+        this.fluentReturnType = fluentReturnType;
         this.response = response;
         this.responseVerificationVariableName = responseVerificationVariableName;
         this.responseVerificationNode = responseVerificationNode;
@@ -41,13 +44,14 @@ public FluentMethodMockUnitTest(
     public FluentMethodMockUnitTest(
             FluentResourceCreateExample fluentResourceCreateExample,
             FluentResourceCollection resourceCollection, FluentCollectionMethod collectionMethod,
-            ProxyMethodExample.Response response,
+            IType fluentReturnType, ProxyMethodExample.Response response,
             String responseVerificationVariableName, ExampleNode responseVerificationNode) {
 
         this.fluentMethodExample = null;
         this.fluentResourceCreateExample = fluentResourceCreateExample;
         this.resourceCollection = resourceCollection;
         this.collectionMethod = collectionMethod;
+        this.fluentReturnType = fluentReturnType;
         this.response = response;
         this.responseVerificationVariableName = responseVerificationVariableName;
         this.responseVerificationNode = responseVerificationNode;
@@ -60,6 +64,13 @@ public FluentMethodExample getFluentMethodExample() {
         return fluentMethodExample;
     }
 
+    /**
+     * @return return type. It could be the inner client model type of Fluent resource create/update/get method.
+     */
+    public IType getFluentReturnType() {
+        return fluentReturnType;
+    }
+
     /**
      * @return example of resource creation, mutually exclusive with {@link #getFluentMethodExample()}.
      */

fluentgen/src/main/java/com/azure/autorest/fluent/model/javamodel/FluentJavaPackage.java

@@ -14,7 +14,7 @@
 import com.azure.autorest.fluent.model.clientmodel.examplemodel.FluentMethodMockUnitTest;
 import com.azure.autorest.fluent.model.projectmodel.Changelog;
 import com.azure.autorest.fluent.model.projectmodel.FluentProject;
-import com.azure.autorest.fluent.template.FluentMethodTestTemplate;
+import com.azure.autorest.fluent.template.FluentMethodMockTestTemplate;
 import com.azure.autorest.fluent.template.FluentLiveTestsTemplate;
 import com.azure.autorest.model.projectmodel.TextFile;
 import com.azure.autorest.fluent.template.ChangelogTemplate;
@@ -116,9 +116,9 @@ public void addOperationUnitTest(FluentMethodMockUnitTest unitTest) {
                 + CodeNamer.toPascalCase(unitTest.getCollectionMethod().getMethodName())
                 + "MockTests";
         JavaFile javaFile = getJavaFileFactory().createTestFile(JavaSettings.getInstance().getPackage("generated"), className);
-        FluentMethodTestTemplate.ClientMethodInfo info = new FluentMethodTestTemplate.ClientMethodInfo(
+        FluentMethodMockTestTemplate.ClientMethodInfo info = new FluentMethodMockTestTemplate.ClientMethodInfo(
                 className, unitTest);
-        FluentMethodTestTemplate.getInstance().write(info, javaFile);
+        FluentMethodMockTestTemplate.getInstance().write(info, javaFile);
         addJavaFile(javaFile);
     }
 

fluentgen/src/main/java/com/azure/autorest/fluent/template/FluentMethodTestTemplate.java → fluentgen/src/main/java/com/azure/autorest/fluent/template/FluentMethodMockTestTemplate.java

@@ -4,10 +4,11 @@
 package com.azure.autorest.fluent.template;
 
 import com.azure.autorest.fluent.model.clientmodel.examplemodel.FluentMethodMockUnitTest;
+import com.azure.autorest.fluent.util.FluentUtils;
 import com.azure.autorest.model.clientmodel.ClassType;
 import com.azure.autorest.model.clientmodel.ClientMethod;
 import com.azure.autorest.model.clientmodel.IType;
-import com.azure.autorest.model.clientmodel.PrimitiveType;
+import com.azure.autorest.model.clientmodel.examplemodel.ExampleHelperFeature;
 import com.azure.autorest.model.clientmodel.examplemodel.ExampleNode;
 import com.azure.autorest.model.javamodel.JavaFile;
 import com.azure.autorest.template.IJavaTemplate;
@@ -30,7 +31,7 @@
 import java.util.HashSet;
 import java.util.Set;
 
-public class FluentMethodTestTemplate implements IJavaTemplate<FluentMethodTestTemplate.ClientMethodInfo, JavaFile> {
+public class FluentMethodMockTestTemplate implements IJavaTemplate<FluentMethodMockTestTemplate.ClientMethodInfo, JavaFile> {
 
     public static class ClientMethodInfo {
         private final String className;
@@ -43,14 +44,14 @@ public ClientMethodInfo(String className, FluentMethodMockUnitTest fluentMethodM
         }
     }
 
-    private static final FluentMethodTestTemplate INSTANCE = new FluentMethodTestTemplate();
+    private static final FluentMethodMockTestTemplate INSTANCE = new FluentMethodMockTestTemplate();
 
     private static final SerializerAdapter SERIALIZER = SerializerFactory.createDefaultManagementSerializerAdapter();
 
-    private FluentMethodTestTemplate() {
+    private FluentMethodMockTestTemplate() {
     }
 
-    public static FluentMethodTestTemplate getInstance() {
+    public static FluentMethodMockTestTemplate getInstance() {
         return INSTANCE;
     }
 
@@ -77,8 +78,12 @@ public void write(ClientMethodInfo info, JavaFile javaFile) {
         String className = info.className;
         FluentMethodMockUnitTest fluentMethodMockUnitTest = info.fluentMethodMockUnitTest;
         ClientMethod clientMethod = fluentMethodMockUnitTest.getCollectionMethod().getInnerClientMethod();
-        final boolean hasReturnValue = clientMethod.getReturnValue().getType() != PrimitiveType.Void;
-        IType fluentReturnType = fluentMethodMockUnitTest.getCollectionMethod().getFluentReturnType();
+        IType fluentReturnType = fluentMethodMockUnitTest.getFluentReturnType();
+        final boolean isResponseType = FluentUtils.isResponseType(fluentReturnType);
+        if (isResponseType) {
+            fluentReturnType = FluentUtils.getValueTypeFromResponseType(fluentReturnType);
+        }
+        final boolean hasReturnValue = fluentReturnType.asNullable() != ClassType.Void;
 
         // method invocation
         String clientMethodInvocationWithResponse;
@@ -92,7 +97,8 @@ public void write(ClientMethodInfo info, JavaFile javaFile) {
         }
         String clientMethodInvocation = exampleMethod.getMethodContent();
         if (hasReturnValue) {
-            clientMethodInvocationWithResponse = fluentReturnType.toString() + " response = " + clientMethodInvocation;
+            // hack on replaceResponseForValue, as in "update" case, "exampleMethod.getMethodContent()" would be a code block, not a single line of code invocation.
+            clientMethodInvocationWithResponse = fluentReturnType + " response = " + (isResponseType ? replaceResponseForValue(clientMethodInvocation) : clientMethodInvocation);
         } else {
             clientMethodInvocationWithResponse = clientMethodInvocation;
         }
@@ -157,9 +163,17 @@ public void write(ClientMethodInfo info, JavaFile javaFile) {
             });
 
             // helper method
-            if (!exampleMethod.getHelperFeatures().isEmpty()) {
+            if (exampleMethod.getHelperFeatures().contains(ExampleHelperFeature.MapOfMethod)) {
                 ModelExampleWriter.writeMapOfMethod(classBlock);
             }
         });
     }
+
+    private static String replaceResponseForValue(String clientMethodInvocation) {
+        if (clientMethodInvocation.endsWith(";")) {
+            clientMethodInvocation = clientMethodInvocation.substring(0, clientMethodInvocation.length() - 1);
+            clientMethodInvocation += ".getValue();";
+        }
+        return clientMethodInvocation;
+    }
 }

fluentgen/src/main/java/com/azure/autorest/fluent/util/FluentUtils.java

@@ -354,14 +354,23 @@ public static boolean exampleIsUpdate(String name) {
         return name.contains("update") && !name.contains("create");
     }
 
-    public static boolean validToGenerateExample(ClientMethod clientMethod) {
+    public static boolean validRequestContentTypeToGenerateExample(ClientMethod clientMethod) {
+        // for now, only accept JSON as request body
+
         String requestContentType = clientMethod.getProxyMethod().getRequestContentType();
         return clientMethod.getProxyMethod().getExamples() != null
                 && requiresExample(clientMethod)
                 // currently only generate for json payload, i.e. "text/json", "application/json"
                 && requestContentType != null && requestContentType.contains("json");
     }
 
+    public static boolean validResponseContentTypeToGenerateExample(ClientMethod clientMethod) {
+        // for now, avoid binary as response body
+
+        IType responseBodyType = clientMethod.getProxyMethod().getResponseBodyType();
+        return !(responseBodyType == ClassType.BinaryData || responseBodyType == GenericType.FluxByteBuffer);
+    }
+
     public static boolean requiresExample(ClientMethod clientMethod) {
         if (clientMethod.getType() == ClientMethodType.SimpleSync
                 || clientMethod.getType() == ClientMethodType.SimpleSyncRestResponse

javagen/src/main/java/com/azure/autorest/model/clientmodel/ProxyMethod.java

@@ -603,7 +603,7 @@ public Builder isResumable(boolean isResumable) {
         /**
          * Sets the metia-types in response.
          *
-         * @param responseContentTypes the metia-types in response
+         * @param responseContentTypes the media-types in response
          * @return the Builder itself
          */
         public Builder responseContentTypes(Set<String> responseContentTypes) {