Test generation for DPG (#2425)

* init test generation code

* add async test generation

* update

* inv reg

* fix

* inv

* optimization

* fix test

* fix ci

* fix pyright

* update

* inv

* black

* merge

* inv

* update

* update

* fix

* mypy fix

* inv

* inv

* inv

* inv

.chronus/changes/test-generation-dev1-2024-3-19-10-27-21.md

@@ -0,0 +1,8 @@
+---
+changeKind: feature
+packages:
+  - "@autorest/python"
+  - "@azure-tools/typespec-python"
+---
+
+add `--generate-test` to generate test for DPG

packages/autorest.python/autorest/codegen/__init__.py

@@ -37,6 +37,7 @@ class OptionsRetriever:
         "multiapi": False,
         "polymorphic-examples": 5,
         "generate-sample": False,
+        "generate-test": False,
         "from-typespec": False,
         "emit-cross-language-definition-file": False,
     }
@@ -332,6 +333,7 @@ def _build_code_model_options(self) -> Dict[str, Any]:
             "packaging_files_config",
             "default_optional_constants_to_none",
             "generate_sample",
+            "generate_test",
             "default_api_version",
             "from_typespec",
             "flavor",
@@ -436,6 +438,7 @@ def get_options(self) -> Dict[str, Any]:
                 "default-optional-constants-to-none"
             ),
             "generate-sample": self._autorestapi.get_boolean_value("generate-sample"),
+            "generate-test": self._autorestapi.get_boolean_value("generate-test"),
             "default-api-version": self._autorestapi.get_value("default-api-version"),
         }
         return {k: v for k, v in options.items() if v is not None}

packages/autorest.python/autorest/codegen/models/code_model.py

@@ -76,6 +76,7 @@ def __init__(
             t for t in self.types_map.values() if isinstance(t, CombinedType) and t.name
         ]
         self.cross_language_package_id = self.yaml_data.get("crossLanguagePackageId")
+        self.for_test: bool = False
 
     @property
     def has_form_data(self) -> bool:

packages/autorest.python/autorest/codegen/models/model_type.py

@@ -148,7 +148,11 @@ def get_json_template_representation(
         if self.discriminated_subtypes:
             # we will instead print the discriminated subtypes
             self._created_json_template_representation = False
-            return self.snake_case_name
+            return (
+                f'"{self.snake_case_name}"'
+                if self.code_model.for_test
+                else self.snake_case_name
+            )
 
         # don't add additional properties, because there's not really a concept of
         # additional properties in the template

packages/autorest.python/autorest/codegen/models/primitive_types.py

@@ -62,7 +62,9 @@ def get_json_template_representation(
             comment = add_to_description(comment, description)
         if comment:
             comment = f"# {comment}"
-        return f"{client_default_value_declaration}{comment}"
+        return client_default_value_declaration + (
+            "" if self.code_model.for_test else comment
+        )
 
     @property
     def default_template_representation_declaration(self) -> str:

packages/autorest.python/autorest/codegen/models/property.py

@@ -134,7 +134,8 @@ def get_json_template_representation(
         description: Optional[str] = None,
     ) -> Any:
         if self.is_multipart_file_input:
-            return "[filetype]" if self.type.type == "list" else "filetype"
+            file_type_str = '"filetype"' if self.code_model.for_test else "filetype"
+            return f"[{file_type_str}]" if self.type.type == "list" else file_type_str
         if self.client_default_value:
             client_default_value_declaration = self.get_declaration(
                 self.client_default_value

packages/autorest.python/autorest/codegen/serializers/__init__.py

@@ -27,6 +27,7 @@
 from .request_builders_serializer import RequestBuildersSerializer
 from .patch_serializer import PatchSerializer
 from .sample_serializer import SampleSerializer
+from .test_serializer import TestSerializer, TestGeneralSerializer
 from .types_serializer import TypesSerializer
 from ..._utils import to_snake_case
 from .._utils import VALID_PACKAGE_MODE
@@ -146,6 +147,14 @@ def _serialize_namespace_level(
         ):
             self._serialize_and_write_sample(env, namespace_path)
 
+        if (
+            self.code_model.options["show_operations"]
+            and self.code_model.has_operations
+            and self.code_model.options["generate_test"]
+            and not self.code_model.options["azure_arm"]
+        ):
+            self._serialize_and_write_test(env, namespace_path)
+
     def serialize(self) -> None:
         env = Environment(
             loader=PackageLoader("autorest.codegen", "templates"),
@@ -631,6 +640,40 @@ def _serialize_and_write_sample(self, env: Environment, namespace_path: Path):
                             log_error = f"error happens in sample {file}: {e}"
                             _LOGGER.error(log_error)
 
+    def _serialize_and_write_test(self, env: Environment, namespace_path: Path):
+        self.code_model.for_test = True
+        out_path = self._package_root_folder(namespace_path) / Path("generated_tests")
+        general_serializer = TestGeneralSerializer(code_model=self.code_model, env=env)
+        self.write_file(
+            out_path / "conftest.py", general_serializer.serialize_conftest()
+        )
+        for is_async in (True, False):
+            async_suffix = "_async" if is_async else ""
+            general_serializer.is_async = is_async
+            self.write_file(
+                out_path / f"testpreparer{async_suffix}.py",
+                general_serializer.serialize_testpreparer(),
+            )
+
+        for client in self.code_model.clients:
+            for og in client.operation_groups:
+                test_serializer = TestSerializer(
+                    self.code_model, env, client=client, operation_group=og
+                )
+                for is_async in (True, False):
+                    try:
+                        test_serializer.is_async = is_async
+                        self.write_file(
+                            out_path
+                            / f"{to_snake_case(test_serializer.test_class_name)}.py",
+                            test_serializer.serialize_test(),
+                        )
+                    except Exception as e:  # pylint: disable=broad-except
+                        # test generation shall not block code generation, so just log error
+                        log_error = f"error happens in test generation for operation group {og.class_name}: {e}"
+                        _LOGGER.error(log_error)
+        self.code_model.for_test = False
+
 
 class JinjaSerializerAutorest(JinjaSerializer, ReaderAndWriterAutorest):
     def __init__(

packages/autorest.python/autorest/codegen/serializers/builder_serializer.py

@@ -4,10 +4,9 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-import json
 from abc import abstractmethod
 from collections import defaultdict
-from typing import Any, Generic, List, Type, TypeVar, Dict, Union, Optional, cast
+from typing import Generic, List, Type, TypeVar, Dict, Union, Optional, cast
 
 from ..models import (
     Operation,
@@ -84,34 +83,13 @@ def _escape_str(input_str: str) -> str:
     return f'"{replace}"'
 
 
-def _improve_json_string(template_representation: str) -> Any:
-    origin = template_representation.split("\n")
-    final = []
-    for line in origin:
-        idx0 = line.find("#")
-        idx1 = line.rfind('"')
-        modified_line = ""
-        if idx0 > -1 and idx1 > -1:
-            modified_line = line[:idx0] + line[idx1:] + "  " + line[idx0:idx1] + "\n"
-        else:
-            modified_line = line + "\n"
-        modified_line = modified_line.replace('"', "").replace("\\", '"')
-        final.append(modified_line)
-    return "".join(final)
-
-
-def _json_dumps_template(template_representation: Any) -> Any:
-    # only for template use, since it wraps everything in strings
-    return _improve_json_string(json.dumps(template_representation, indent=4))
-
-
 def _get_polymorphic_subtype_template(polymorphic_subtype: ModelType) -> List[str]:
     retval: List[str] = []
     retval.append("")
     retval.append(
         f'# JSON input template for discriminator value "{polymorphic_subtype.discriminator_value}":'
     )
-    subtype_template = _json_dumps_template(
+    subtype_template = utils.json_dumps_template(
         polymorphic_subtype.get_json_template_representation(),
     )
 
@@ -229,7 +207,7 @@ def _get_json_response_template_to_status_codes(
         if not json_template:
             continue
         status_codes = [str(status_code) for status_code in response.status_codes]
-        response_json = _json_dumps_template(json_template)
+        response_json = utils.json_dumps_template(json_template)
         retval[response_json].extend(status_codes)
     return retval
 
@@ -440,7 +418,7 @@ def _json_input_example_template(self, builder: BuilderType) -> List[str]:
         template.append(
             "# JSON input template you can fill out and use as your body input."
         )
-        json_template = _json_dumps_template(
+        json_template = utils.json_dumps_template(
             json_type.get_json_template_representation(),
         )
         template.extend(