Merge remote-tracking branch 'origin/dev' into vm-diagnostics-migration

Author: william051200

src/azure-cli-core/azure/cli/core/__init__.py

@@ -1134,22 +1134,17 @@ def _load_command_loader(loader, args, name, prefix):
             logger.debug("Module '%s' is missing `get_command_loader` entry.", name)
 
     command_table = {}
+    command_loader = None
 
     if loader_cls:
         command_loader = loader_cls(cli_ctx=loader.cli_ctx)
-        loader.loaders.append(command_loader)  # This will be used by interactive
         if command_loader.supported_resource_type():
             command_table = command_loader.load_command_table(args)
-            if command_table:
-                for cmd in list(command_table.keys()):
-                    # TODO: If desired to for extension to patch module, this can be uncommented
-                    # if loader.cmd_to_loader_map.get(cmd):
-                    #    loader.cmd_to_loader_map[cmd].append(command_loader)
-                    # else:
-                    loader.cmd_to_loader_map[cmd] = [command_loader]
     else:
         logger.debug("Module '%s' is missing `COMMAND_LOADER_CLS` entry.", name)
-    return command_table, command_loader.command_group_table
+
+    group_table = command_loader.command_group_table if command_loader else {}
+    return command_table, group_table, command_loader
 
 
 def _load_extension_command_loader(loader, args, ext):

src/azure-cli-core/azure/cli/core/commands/__init__.py

@@ -230,7 +230,7 @@ def load_command_table(self, args):
             if command_table:
                 module_command_table.update(command_table)
                 loader.loaders.append(command_loader)  # this will be used later by the load_arguments method
-        return module_command_table, command_loader.command_group_table
+        return module_command_table, command_loader.command_group_table, command_loader
 
     expected_command_index = {'hello': ['azure.cli.command_modules.hello', 'azext_hello2', 'azext_hello1'],
                               'extra': ['azure.cli.command_modules.extra']}
@@ -260,6 +260,41 @@ def test_register_command_from_extension(self):
         self.assertTrue(isinstance(hello_overridden_cmd.command_source, ExtensionCommandSource))
         self.assertTrue(hello_overridden_cmd.command_source.overrides_command)
 
+    @mock.patch('importlib.import_module', _mock_import_lib)
+    @mock.patch('pkgutil.iter_modules', _mock_iter_modules)
+    @mock.patch('azure.cli.core.commands._load_command_loader', _mock_load_command_loader)
+    @mock.patch('azure.cli.core.extension.get_extension_modname', _mock_get_extension_modname)
+    @mock.patch('azure.cli.core.extension.get_extensions', _mock_get_extensions)
+    def test_cmd_to_loader_map_populated_after_parallel_loading(self):
+        """
+        Validates that all commands in command_table have corresponding entries in cmd_to_loader_map.
+        """
+        cli = DummyCli()
+        loader = cli.commands_loader
+
+        # Load all commands (triggers parallel module loading)
+        cmd_tbl = loader.load_command_table(None)
+        
+        # Verify EVERY command in command_table has an entry in cmd_to_loader_map
+        # This is exactly what azdev does before it hits KeyError
+        for cmd_name in cmd_tbl:
+            # This should NOT raise KeyError
+            self.assertIn(cmd_name, loader.cmd_to_loader_map, 
+                         f"Command '{cmd_name}' missing from cmd_to_loader_map - "
+                         f"would cause KeyError in azdev command-change meta-export")
+            
+            # Verify the entry is a list with at least one loader
+            loaders = loader.cmd_to_loader_map[cmd_name]
+            self.assertIsInstance(loaders, list, 
+                                 f"cmd_to_loader_map['{cmd_name}'] should be a list")
+            self.assertGreater(len(loaders), 0, 
+                              f"cmd_to_loader_map['{cmd_name}'] should have at least one loader")
+        
+        # Verify all expected commands are present
+        expected_commands = {'hello mod-only', 'hello overridden', 'extra final', 'hello ext-only'}
+        actual_commands = set(cmd_tbl.keys())
+        self.assertEqual(expected_commands, actual_commands)
+
     @mock.patch('importlib.import_module', _mock_import_lib)
     @mock.patch('pkgutil.iter_modules', _mock_iter_modules)
     @mock.patch('azure.cli.core.commands._load_command_loader', _mock_load_command_loader)
@@ -432,12 +467,12 @@ def test_command_index_positional_argument(self):
         # Test command index is built for command with positional argument
         cmd_tbl = loader.load_command_table(["extra", "extra", "positional_argument"])
         self.assertDictEqual(INDEX[CommandIndex._COMMAND_INDEX], self.expected_command_index)
-        self.assertEqual(list(cmd_tbl), ['hello mod-only', 'hello overridden', 'extra final', 'hello ext-only'])
+        self.assertSetEqual(set(cmd_tbl), {'hello mod-only', 'hello overridden', 'extra final', 'hello ext-only'})
 
         # Test command index is used by command with positional argument
         cmd_tbl = loader.load_command_table(["hello", "mod-only", "positional_argument"])
         self.assertDictEqual(INDEX[CommandIndex._COMMAND_INDEX], self.expected_command_index)
-        self.assertEqual(list(cmd_tbl), ['hello mod-only', 'hello overridden', 'hello ext-only'])
+        self.assertSetEqual(set(cmd_tbl), {'hello mod-only', 'hello overridden', 'hello ext-only'})
 
         # Test command index is used by command with positional argument
         cmd_tbl = loader.load_command_table(["extra", "final", "positional_argument2"])

src/azure-cli-core/azure/cli/core/tests/test_command_registration.py

@@ -0,0 +1,57 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import unittest
+
+from azure.cli.core.mock import DummyCli
+from azure.cli.core import MainCommandsLoader
+
+
+class CommandTableIntegrityTest(unittest.TestCase):
+
+    def setUp(self):
+        self.cli_ctx = DummyCli()
+
+    def test_command_table_integrity(self):
+        """Test command table loading produces valid, complete results."""
+
+        # Load command table using current implementation
+        loader = MainCommandsLoader(self.cli_ctx)
+        loader.load_command_table([])
+
+        # Test invariants that should always hold:
+
+        # 1. No corruption/duplicates
+        command_names = list(loader.command_table.keys())
+        unique_command_names = set(command_names)
+        self.assertEqual(len(unique_command_names), len(command_names), "No duplicate commands")
+
+        # 2. Core functionality exists (high-level groups that should always exist)
+        core_groups = ['vm', 'network', 'resource', 'account', 'group']
+        existing_groups = {cmd.split()[0] for cmd in loader.command_table.keys() if ' ' in cmd}
+        missing_core = [group for group in core_groups if group not in existing_groups]
+        self.assertEqual(len(missing_core), 0, f"Missing core command groups: {missing_core}")
+
+        # 3. Structural integrity
+        commands_without_source = []
+        for cmd_name, cmd_obj in loader.command_table.items():
+            if not hasattr(cmd_obj, 'command_source') or not cmd_obj.command_source:
+                commands_without_source.append(cmd_name)
+
+        self.assertEqual(len(commands_without_source), 0,
+                         f"Commands missing source: {commands_without_source[:5]}...")
+
+        # 4. Basic sanity - we loaded SOMETHING
+        self.assertGreater(len(loader.command_table), 0, "Commands were loaded")
+        self.assertGreater(len(loader.command_group_table), 0, "Groups were loaded")
+
+        # 5. Verify core groups are properly represented
+        found_core_groups = sorted(existing_groups & set(core_groups))
+        self.assertGreaterEqual(len(found_core_groups), 3,
+                                f"At least 3 core command groups should be present, found: {found_core_groups}")
+
+
+if __name__ == '__main__':
+    unittest.main()

src/azure-cli-core/azure/cli/core/tests/test_command_table_integrity.py

@@ -188,7 +188,7 @@ def load_command_table(self, args):
             if command_table:
                 module_command_table.update(command_table)
                 loader.loaders.append(command_loader)  # this will be used later by the load_arguments method
-        return module_command_table, command_loader.command_group_table
+        return module_command_table, command_loader.command_group_table, command_loader
 
     @mock.patch('importlib.import_module', _mock_import_lib)
     @mock.patch('pkgutil.iter_modules', _mock_iter_modules)

src/azure-cli-core/azure/cli/core/tests/test_parser.py

@@ -185,6 +185,28 @@
     "ussecwest": "ussecwest",
 }
 
+# mapping for azure bleu cloud
+AzureBleuLocationToOmsRegionCodeMap = {
+    "bleufrancecentral": "BLEUC",
+    "bleufrancesouth": "BLEUS",
+}
+
+AzureBleuRegionToOmsRegionMap = {
+    "bleufrancecentral": "bleufrancecentral",
+    "bleufrancesouth": "bleufrancesouth",
+}
+
+# mapping for azure delos cloud
+AzureDelosLocationToOmsRegionCodeMap = {
+    "deloscloudgermanycentral": "DELOSC",
+    "deloscloudgermanynorth": "DELOSN",
+}
+
+AzureDelosRegionToOmsRegionMap = {
+    "deloscloudgermanycentral": "deloscloudgermanycentral",
+    "deloscloudgermanynorth": "deloscloudgermanynorth",
+}
+
 ContainerInsightsStreams = [
     "Microsoft-ContainerLog",
     "Microsoft-ContainerLogV2-HighScale",
@@ -245,7 +267,20 @@ def ensure_default_log_analytics_workspace_for_monitoring(
         workspace_region_code = AzureUSSecLocationToOmsRegionCodeMap.get(
             workspace_region, "USSW"
         )
-
+    elif cloud_name.lower() == "azurebleucloud":
+        workspace_region = AzureBleuRegionToOmsRegionMap.get(
+            rg_location, "bleufrancecentral"
+        )
+        workspace_region_code = AzureBleuLocationToOmsRegionCodeMap.get(
+            workspace_region, "BLEUC"
+        )
+    elif cloud_name.lower() == "azuredeloscloud":
+        workspace_region = AzureDelosRegionToOmsRegionMap.get(
+            rg_location, "deloscloudgermanycentral"
+        )
+        workspace_region_code = AzureDelosLocationToOmsRegionCodeMap.get(
+            workspace_region, "DELOSC"
+        )
     else:
         logger.error(
             "AKS Monitoring addon not supported in cloud : %s", cloud_name