Fix data-plane spec discovery and code generation

src/aaz_dev/app/tests/common.py

@@ -27,4 +27,10 @@ def cleanup_dev_folder(self):
 
     def setUp(self):
         os.makedirs(self.AAZ_FOLDER, exist_ok=True)
+        # the fake aaz repo is wiped for every test, seed it with an empty command tree so that
+        # loading the command tree doesn't fail on a missing readme.md
+        commands_folder = os.path.join(self.AAZ_FOLDER, 'Commands')
+        os.makedirs(commands_folder, exist_ok=True)
+        with open(os.path.join(commands_folder, 'readme.md'), 'w', encoding='utf-8') as f:
+            f.write("# Atomic Azure CLI Commands\n\n## Groups\n\n")
 

src/aaz_dev/cli/api/_cmds.py

@@ -230,13 +230,16 @@ def _collect_resources(spec):
 
             rp = module_manager.get_openapi_resource_provider(r.name)
             tag = r.default_tag
-
+            if not tag:
+                raise InvalidAPIUsage(f"Cannot generate `{spec}`: resource provider `{r.name}` has no default tag.")
             resource_map = rp.get_resource_map_by_tag(tag)
             if not resource_map:
                 raise InvalidAPIUsage(f"Tag `{tag}` is not exist.")
 
             results[rp.name] = (resource_map, tag)
 
+        if not results:
+            raise InvalidAPIUsage(f"Cannot generate `{spec}`: no OpenAPI resources were selected.")
         return results
 
     def _normalize_resource_map(resource_map, tag):
@@ -311,8 +314,8 @@ def to_cli():
                 v = v_list[0]
                 cfg_reader = AAZSpecsManager().load_resource_cfg_reader(Config.DEFAULT_PLANE, resource_id, v)
                 if not cfg_reader:
-                    logger.error(f"Command models not exist in aaz for resource: {resource_id} version: {v}.")
-                    continue
+                    raise InvalidAPIUsage(
+                        f"Command models not exist in aaz for resource: {resource_id} version: {v}.")
 
                 for cmd_names, command in cfg_reader.iter_commands():
                     key = tuple(cmd_names)
@@ -321,6 +324,9 @@ def to_cli():
 
                     commands_map[key] = command.version
 
+        if not commands_map:
+            raise InvalidAPIUsage(f"Cannot generate `{spec}`: no commands were generated.")
+
         if cli_path is not None:
             assert Config.CLI_PATH is not None
             manager = AzMainManager()

src/aaz_dev/cli/controller/az_module_manager.py

@@ -90,7 +90,47 @@ def update_module(self, mod_name, profiles, **kwargs):
         return module
 
     _def_load_command_table = re.compile(r"^(\s+)def\s+load_command_table\(\s*self,\s+(\w+)\s*\):(.*)?$")
-    _def_import_load_aaz = re.compile(r"\s+(import\s+(\w+.)*load_aaz_command_table)\s*$")
+
+    @staticmethod
+    def _loads_aaz_commands(module, line):
+        owner = next((
+            node for node in ast.walk(module) if isinstance(node, ast.ClassDef)
+            and any(isinstance(member, ast.FunctionDef) and member.lineno == line for member in node.body)
+        ), None)
+        if owner is None:
+            return False
+        methods = {node.name: node for node in owner.body if isinstance(node, ast.FunctionDef)}
+        loaders = {"load_aaz_command_table", "load_aaz_command_table_args_guided"}
+        pending = ["load_command_table"]
+        visited = set()
+        while pending:
+            name = pending.pop()
+            if name in visited or name not in methods:
+                continue
+            visited.add(name)
+            nodes = []
+            stack = list(methods[name].body)
+            while stack:
+                node = stack.pop()
+                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)):
+                    continue
+                nodes.append(node)
+                stack.extend(ast.iter_child_nodes(node))
+            aliases = loaders | {
+                alias.asname or alias.name
+                for node in [*module.body, *nodes]
+                if isinstance(node, ast.ImportFrom) and node.module == "azure.cli.core.aaz"
+                for alias in node.names if alias.name in loaders
+            }
+            for node in nodes:
+                if not isinstance(node, ast.Call):
+                    continue
+                if isinstance(node.func, ast.Name) and node.func.id in aliases:
+                    return True
+                if (isinstance(node.func, ast.Attribute) and isinstance(node.func.value, ast.Name)
+                        and node.func.value.id == "self"):
+                    pending.append(node.func.attr)
+        return False
 
     def _patch_module(self, mod_name):
         """Patch the __init__.py file of module"""
@@ -107,10 +147,6 @@ def _patch_module(self, mod_name):
         space = None
         for idx in range(len(lines)):
             line = lines[idx]
-            if self._def_import_load_aaz.findall(line):
-                # already patched
-                logger.debug(f"Module is already patched")
-                return
             if start_line is None:
                 def_match = self._def_load_command_table.match(line)
                 if def_match:
@@ -131,6 +167,14 @@ def _patch_module(self, mod_name):
         if start_line is None:
             raise exceptions.InvalidAPIUsage(f"Patch Module failed: Cannot find load_command_table function in file: {file}")
 
+        try:
+            module = ast.parse('\n'.join(lines), filename=file)
+        except SyntaxError as err:
+            raise exceptions.InvalidAPIUsage(f"Patch Module failed: Invalid Python in file: {file}: {err}") from err
+        if self._loads_aaz_commands(module, start_line + 1):
+            logger.debug("Module is already patched")
+            return
+
         insert_lines = [
             f"{space}{space}from azure.cli.core.aaz import load_aaz_command_table",
             f"{space}{space}try:",

src/aaz_dev/cli/tests/test_codegen_regressions.py

@@ -0,0 +1,147 @@
+import inspect
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from types import SimpleNamespace
+from unittest import TestCase
+from unittest.mock import patch
+
+from cli.api import _cmds
+from cli.controller.az_module_manager import AzModuleManager
+from swagger.model.specs import SwaggerSpecs, TypeSpecResourceProvider
+from utils.config import Config
+from utils.plane import PlaneEnum
+
+
+class CodegenRegressionTest(TestCase):
+    def test_generate_rejects_incomplete_selections_before_updating_cli(self):
+        resources = {"/test": {"v1": object()}}
+        good = SimpleNamespace(name="Good", default_tag="v1", get_resource_map_by_tag=lambda _: resources)
+        missing = SimpleNamespace(name="Missing", default_tag=None)
+        typespec = TypeSpecResourceProvider("Test", [], None)
+        reader = SimpleNamespace(iter_commands=lambda: iter([
+            (["test", "show"], SimpleNamespace(version="v1")),
+        ]))
+        empty_reader = SimpleNamespace(iter_commands=lambda: iter([]))
+
+        cases = [
+            ("no providers", [], reader, False),
+            ("typespec only", [typespec], reader, False),
+            ("missing tag", [missing], reader, False),
+            ("partial selection", [good, missing], reader, False),
+            ("missing models", [good], None, False),
+            ("no commands", [good], empty_reader, False),
+            ("valid selection", [good], reader, True),
+        ]
+        with TemporaryDirectory() as folder:
+            (Path(folder) / "specification" / "test").mkdir(parents=True)
+            for name, providers, cfg_reader, valid in cases:
+                with self.subTest(name=name):
+                    module_manager = SimpleNamespace(
+                        get_resource_providers=lambda: providers,
+                        get_openapi_resource_provider=lambda name: next(r for r in providers if r.name == name),
+                    )
+                    discovery = SimpleNamespace(get_module_manager=lambda *args: module_manager)
+                    with patch.object(Config, "SWAGGER_PATH", folder), \
+                            patch.object(Config, "CLI_EXTENSION_PATH", folder), \
+                            patch.object(_cmds, "SwaggerSpecsManager", return_value=discovery), \
+                            patch.object(_cmds, "WorkspaceManager") as workspaces, \
+                            patch.object(_cmds, "AAZSpecsManager") as aaz, \
+                            patch.object(_cmds, "AzExtensionManager") as extensions:
+                        workspaces.new.return_value.is_in_memory = True
+                        workspaces.new.return_value.iter_command_tree_nodes.return_value = []
+                        workspaces.new.return_value.iter_command_tree_leaves.return_value = []
+                        aaz.return_value.load_resource_cfg_reader.return_value = cfg_reader
+                        manager = extensions.return_value
+                        manager.has_module.return_value = True
+                        manager.load_module.return_value = SimpleNamespace(profiles={})
+                        if valid:
+                            inspect.unwrap(_cmds.generate.callback)("test", "test")
+                            manager.update_module.assert_called_once()
+                            profile = manager.update_module.call_args.args[1]["latest"]
+                            self.assertEqual(profile.command_groups["test"].commands["show"].version, "v1")
+                        else:
+                            with self.assertRaises(SystemExit) as error:
+                                inspect.unwrap(_cmds.generate.callback)("test", "test")
+                            self.assertEqual(error.exception.code, 1)
+                            manager.update_module.assert_not_called()
+                            manager.create_new_mod.assert_not_called()
+                            if name in ("no providers", "typespec only", "missing tag", "partial selection"):
+                                workspaces.new.assert_not_called()
+
+    def test_submodules_preserve_typespec_and_openapi_layouts(self):
+        for plane, segment in ((PlaneEnum.Mgmt, "resource-manager"), (PlaneEnum._Data, "data-plane")):
+            with self.subTest(plane=plane), TemporaryDirectory() as folder:
+                root = Path(folder) / "specification" / "test"
+                entry = root / "Service.TypeSpec"
+                entry.mkdir(parents=True)
+                (entry / "main.tsp").write_text(
+                    ("@armProviderNamespace\n" if plane == PlaneEnum.Mgmt else "") +
+                    "namespace Test.Service;\n", encoding="utf-8",
+                )
+                (entry / "tspconfig.yaml").write_text("{}\n", encoding="utf-8")
+                specs = SwaggerSpecs(folder)
+                get_module = specs.get_mgmt_plane_module if plane == PlaneEnum.Mgmt else specs.get_data_plane_module
+                with patch.object(Config, "SWAGGER_PATH", folder):
+                    self.assertEqual(Path(get_module("test", plane=plane).folder_path), root)
+                    module = get_module("test", entry.name, plane=plane)
+                    self.assertIsNotNone(module)
+                    self.assertEqual(Path(module.folder_path), entry)
+                    providers = module.get_resource_providers()
+                    self.assertEqual([r.name for r in providers], ["Test.Service"])
+                    self.assertEqual(len(providers[0].entry_files), 1)
+
+                    (root / segment / entry.name).mkdir(parents=True)
+                    module = get_module("test", entry.name, plane=plane)
+                    self.assertEqual(Path(module.folder_path), entry)
+                    self.assertEqual([r.name for r in module.get_resource_providers()], ["Test.Service"])
+
+                    nested = root / segment / "group" / "nested"
+                    (nested / "Test.Provider" / "stable").mkdir(parents=True)
+                    (root / "group").mkdir()
+                    module = get_module("test", "group", "nested", plane=plane)
+                    self.assertEqual(Path(module.folder_path), nested)
+                    self.assertEqual(module.names, ["test", "group", "nested"])
+                    self.assertEqual([r.name for r in module.get_resource_providers()], ["Test.Provider"])
+                    self.assertIsNone(get_module("test", "missing", plane=plane))
+
+    def test_patch_checks_called_loaders_not_import_text(self):
+        plain = "class Loader:\n    def load_command_table(self, args):\n        return {}\n"
+        helper = (
+            "    def _load(self, args):\n"
+            "        from azure.cli.core.aaz import {loader}\n"
+            "        {loader}(self, 'test.aaz', args)\n"
+        )
+        cases = [
+            ("plain", plain, True),
+            ("comment", plain + "    # from azure.cli.core.aaz import load_aaz_command_table\n", True),
+            ("unused import", "from azure.cli.core.aaz import load_aaz_command_table\n" + plain, True),
+            ("unused helper", plain + helper.format(loader="load_aaz_command_table"), True),
+            ("nested unused helper", plain.replace(
+                "        return {}",
+                "        def unused():\n"
+                "            from azure.cli.core.aaz import load_aaz_command_table\n"
+                "            load_aaz_command_table(self, 'test.aaz', args)\n"
+                "        return {}",
+            ), True),
+        ]
+        for loader in ("load_aaz_command_table", "load_aaz_command_table_args_guided"):
+            cases.append((loader, plain.replace("return {}", "return self._load(args)") +
+                          helper.format(loader=loader), False))
+            cases.append((loader + " alias",
+                          "from azure.cli.core.aaz import " + loader + " as load_commands\n" +
+                          plain.replace("return {}", "return load_commands(self, 'test.aaz', args)"), False))
+        cases.append(("helper cycle", plain.replace("return {}", "return self._load(args)") +
+                      "    def _load(self, args):\n        return self.load_command_table(args)\n", True))
+        with TemporaryDirectory() as folder:
+            init = Path(folder) / "__init__.py"
+            manager = AzModuleManager()
+            manager.get_aaz_path = lambda _: str(Path(folder) / "aaz")
+            for name, source, needs_patch in cases:
+                with self.subTest(name=name):
+                    init.write_text(source, encoding="utf-8")
+                    patches = list(manager._patch_module("test"))
+                    self.assertEqual(len(patches), int(needs_patch))
+                    if patches:
+                        compile(patches[0][1], str(init), "exec")
+                        init.write_text(patches[0][1], encoding="utf-8")
+                        self.assertEqual(list(manager._patch_module("test")), [])

src/aaz_dev/command/tests/api_tests/test_editor.py

@@ -2471,7 +2471,7 @@ def test_dataplane_workspace(self, ws_name):
     @workspace_name("test_dataplane_monitor_metrics")
     def test_dataplane_monitor_metrics(self, ws_name):
         module = "monitor"
-        resource_provider = "Microsoft.Insights"
+        resource_provider = "Insights"
         api_version = '2023-05-01-preview'
         with self.app.test_client() as c:
             rv = c.post(f"/AAZ/Editor/Workspaces", json={
@@ -2643,7 +2643,7 @@ def test_mgmt_attestation(self, ws_name):
     @workspace_name("test_dataplane_attestation")
     def test_dataplane_attestation(self, ws_name):
         module = "attestation"
-        resource_provider = "Microsoft.Attestation"
+        resource_provider = "Attestation"
         api_version = '2022-09-01-preview'
 
         with self.app.test_client() as c:
@@ -2693,6 +2693,7 @@ def test_dataplane_attestation(self, ws_name):
                         'options': ['provider-name'],
                         'required': True,
                         'type': 'string',
+                        'format': {'pattern': '^[a-zA-Z0-9-]{3,24}$'},
                         'group': 'Client',
                         'idPart': 'name',
                         'help': {'short': 'Name of the attestation provider.'},

src/aaz_dev/swagger/api/specs.py

@@ -247,7 +247,7 @@ def get_resource_in_module(plane, mod_names, resource_id):
 )
 def get_resource_version_in_rp(plane, mod_names, rp_name, resource_id, version):
     specs_module_manager = SwaggerSpecsManager().get_module_manager(plane, mod_names)
-    resource = specs_module_manager.get_resource_in_version(rp_name, resource_id, version)
+    resource = specs_module_manager.get_resource_in_version(resource_id, version, rp_name=rp_name)
     result = {
         "url": url_for('swagger.get_resource_version_in_rp',
                        plane=plane, mod_names=mod_names, rp_name=resource.rp_name,

src/aaz_dev/swagger/model/specs/_resource_provider.py

@@ -3,7 +3,6 @@
 import logging
 import os
 import re
-import sys
 from collections import OrderedDict
 
 import yaml
@@ -91,15 +90,17 @@ def tags(self):
     @property
     def default_tag(self):
         if self._default_tag is None:
+            if not self._readme_paths:
+                return None
             with open(self._readme_paths[0], "r", encoding="utf-8") as f:
                 content = f.read()
 
             try:
                 self._default_tag = re.findall(r"tag:\s*(.+)", content)[0]
 
             except IndexError:
-                logger.error(f"Cannot find default tag in resource provider: {self.name}.", exc_info=True)
-                raise sys.exit(1)
+                logger.error(f"Cannot find default tag in resource provider: {self.name}.")
+                return None
 
         return self._default_tag
 

src/aaz_dev/swagger/model/specs/_swagger_module.py

@@ -118,16 +118,21 @@ def _get_openapi_resource_providers(self):
             return rp
         for name in os.listdir(folder_path):
             path = os.path.join(folder_path, name)
-            if os.path.isdir(path):
-                if name.lower() in ('preview', 'stable'):
-                    continue
-                name_parts = name.split('.')
-                if len(name_parts) >= 2:
-                    readme_paths = [*_search_readme_md_paths(path, search_parent=True)]
-                    rp.append(OpenAPIResourceProvider(name, path, readme_paths, swagger_module=self))
-                elif name.lower() != 'common':
-                    sub_module = DataPlaneModule(plane=self.plane, name=name, folder_path=path, parent=self)
-                    rp.extend(sub_module.get_resource_providers())
+            if not os.path.isdir(path):
+                continue
+            if name.startswith('.') or name.lower() in ('preview', 'stable', 'common', 'examples'):
+                continue
+            # A data-plane resource provider folder is either named `Foo.Bar` (legacy) or holds the
+            # `stable`/`preview` version folders directly (unified folder structure). Anything else
+            # is a grouping folder to descend into.
+            if '.' in name or _has_version_folder(path):
+                readme_paths = [*_search_readme_md_paths(path, search_parent=True)]
+                rp.append(OpenAPIResourceProvider(name, path, readme_paths, swagger_module=self))
+            else:
+                sub_module = DataPlaneModule(plane=self.plane, name=name, folder_path=path, parent=self)
+                # only descend for openapi: typespec entry files are found by walking the whole
+                # module folder already, descending again would report every provider twice.
+                rp.extend(sub_module._get_openapi_resource_providers())
         return rp
 
     def _get_typespec_resource_providers(self):
@@ -145,6 +150,10 @@ def _get_typespec_resource_providers(self):
         return [*rp.values()]
 
 
+def _has_version_folder(path):
+    return any(os.path.isdir(os.path.join(path, name)) for name in ('stable', 'preview'))
+
+
 def _search_readme_md_paths(path, search_parent=False):
     # Check parent directory first if requested
     if search_parent: