Merge remote-tracking branch 'origin' into vm-unmanaged-disk-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

@@ -843,19 +843,29 @@ def flexible_server_restore(cmd, client, resource_group_name, server_name, sourc
         else:
             parameters.network = source_server_object.network
 
-    except Exception as e:
-        raise ResourceNotFoundError(e)
-
-    resolve_poller(
-        client.begin_create(resource_group_name, server_name, parameters), cmd.cli_ctx,
-        'Restore Server')
+    except HttpResponseError as exc:
+        raise ResourceNotFoundError(exc) from exc
+
+    def _begin_network_update():
+        restore_server_object = client.get(resource_group_name, server_name)
+        restore_server_network = restore_server_object.network
+        restore_server_network.public_network_access = public_access if public_access else source_server_object.network.public_network_access
+        update_parameter = models.ServerForUpdate(network=restore_server_network)
+        return client.begin_update(resource_group_name, server_name, update_parameter)
+
+    create_poller = sdk_no_wait(no_wait, client.begin_create, resource_group_name, server_name, parameters)
+    if no_wait:
+        def _post_create_update(poller):
+            try:
+                _begin_network_update()
+            except (HttpResponseError, CLIError) as ex:
+                logger.warning('Skipping post-restore network update: %s', ex)
 
-    restore_server_object = client.get(resource_group_name, server_name)
-    restore_server_network = restore_server_object.network
-    restore_server_network.public_network_access = public_access if public_access else source_server_object.network.public_network_access
-    update_parameter = models.ServerForUpdate(network=restore_server_network)
+        create_poller.add_done_callback(_post_create_update)
+        return create_poller
 
-    return sdk_no_wait(no_wait, client.begin_update, resource_group_name, server_name, update_parameter)
+    resolve_poller(create_poller, cmd.cli_ctx, 'Restore Server')
+    return sdk_no_wait(no_wait, _begin_network_update)
 
 
 # pylint: disable=too-many-locals, too-many-statements, raise-missing-from

src/azure-cli/azure/cli/command_modules/mysql/custom.py

@@ -120,6 +120,12 @@ def test_mysql_flexible_server_georestore_mgmt(self, resource_group):
     def test_mysql_flexible_server_georestore_update_mgmt(self, resource_group):
         self._test_flexible_server_georestore_update_mgmt('mysql', resource_group)
 
+    @AllowLargeResponse()
+    @ResourceGroupPreparer(location=DEFAULT_LOCATION)
+    @live_only()
+    def test_mysql_flexible_server_restore_no_wait_mgmt(self, resource_group):
+        self._test_flexible_server_restore_no_wait('mysql', resource_group)
+
     @AllowLargeResponse()
     @ResourceGroupPreparer(location=DEFAULT_LOCATION)
     def test_mysql_flexible_server_gtid_reset(self, resource_group):
@@ -698,6 +704,23 @@ def _test_flexible_server_georestore_update_mgmt(self, database_engine, resource
         self.cmd('{} flexible-server delete -g {} -n {} --yes'.format(database_engine, resource_group, source_server))
         self.cmd('{} flexible-server delete -g {} -n {} --yes'.format(database_engine, resource_group, target_server))
 
+    def _test_flexible_server_restore_no_wait(self, database_engine, resource_group):
+        location = DEFAULT_LOCATION
+        source_server = self.create_random_name(SERVER_NAME_PREFIX, SERVER_NAME_MAX_LENGTH)
+        target_server = self.create_random_name(SERVER_NAME_PREFIX, SERVER_NAME_MAX_LENGTH)
+
+        self.cmd('{} flexible-server create -g {} -n {} -l {} --public-access None --tier GeneralPurpose --sku-name {}'
+                 .format(database_engine, resource_group, source_server, location, DEFAULT_GENERAL_PURPOSE_SKU))
+
+        self.cmd('{} flexible-server restore -g {} --name {} --source-server {} --no-wait'
+                 .format(database_engine, resource_group, target_server, source_server))
+
+        self.cmd('{} flexible-server wait -g {} -n {} --created --interval 30 --timeout 600'
+                 .format(database_engine, resource_group, target_server))
+
+        self.cmd('{} flexible-server delete -g {} -n {} --yes'.format(database_engine, resource_group, source_server))
+        self.cmd('{} flexible-server delete -g {} -n {} --yes'.format(database_engine, resource_group, target_server))
+
 
     def _test_flexible_server_byok_mgmt(self, database_engine, resource_group, vault_name, backup_vault_name=None):
         key_name = self.create_random_name('rdbmskey', 32)

src/azure-cli/azure/cli/command_modules/mysql/tests/latest/recordings/test_mysql_flexible_server_restore_no_wait_mgmt.yaml

@@ -119,10 +119,18 @@ def _get_client(client, kwargs):
         dir_client = client.get_directory_client(directory_path=total_path)
         exists = False
         from azure.core.exceptions import ClientAuthenticationError
+        from azure.core.exceptions import ResourceExistsError
         try:
             exists = dir_client.exists()
         except ClientAuthenticationError:
             exists = False
+        except ResourceExistsError as ex:
+            if hasattr(ex, "error_code") and ex.error_code == "DeletePending":
+                # translate delete pending flag as file/dir not exists.
+                exists = False
+            else:
+                raise ex
+
         if not exists:
             dir_client = client.get_directory_client(directory_path=directory_path)
             client = dir_client.get_file_client(file_name=file_name)

src/azure-cli/azure/cli/command_modules/mysql/tests/latest/test_mysql_scenario.py

@@ -1452,6 +1452,15 @@ def _validate_generation_version_and_trusted_launch(cmd, namespace):
                                                        namespace.os_offer, namespace.os_sku)
             vm_image_info = client.get(namespace.location, namespace.os_publisher, namespace.os_offer,
                                        namespace.os_sku, os_version)
+
+            if vm_image_info.image_deprecation_status.image_state == 'ScheduledForDeprecation':
+                logger.warning(
+                    'Warning: This image %s is scheduled for deprecation and will be blocked after %s.\n'
+                    'VM / VMSS creation is allowed temporarily, but future deployments, redeployments, or '
+                    'scale‑out operations may fail after this date.\n'
+                    'Consider switching to a supported image now.', namespace.image,
+                    vm_image_info.image_deprecation_status.scheduled_deprecation_time.strftime("%B %d, %Y"))
+
             generation_version = vm_image_info.hyper_v_generation if hasattr(vm_image_info,
                                                                              'hyper_v_generation') else None
             features = vm_image_info.features if hasattr(vm_image_info, 'features') and vm_image_info.features else []