diff --git a/src/aks-preview/HISTORY.rst b/src/aks-preview/HISTORY.rst index be6aa0cb576..3494789c8a4 100644 --- a/src/aks-preview/HISTORY.rst +++ b/src/aks-preview/HISTORY.rst @@ -11,6 +11,7 @@ To release a new version, please select a new version number (usually plus 1 to Pending +++++++ +* `az aks create/update`: Add `--enable-azure-managed-lustre` to install the Azure Managed Lustre cluster extension at version 0.6.0 on the stable release train with automatic upgrades disabled. Add `az aks update --disable-azure-managed-lustre` to uninstall it. * `az aks nodepool update`: Preserve the existing GPU management mode when `--enable-managed-gpu` is omitted, including when enabling, updating, or disabling the cluster autoscaler. * `az aks alert-config add`: Reject an empty `--name` before looking up existing configurations instead of reporting that it already exists. * `az aks nodepool scale`: add `--use-patch-api` to optionally scale a VMSS node pool via the new dedicated PATCH agent pool API (scales to the target count without triggering full reconciliation). The default behavior continues to use the PUT agent pool API. diff --git a/src/aks-preview/README.rst b/src/aks-preview/README.rst index a1674f3c1c2..c142290a990 100644 --- a/src/aks-preview/README.rst +++ b/src/aks-preview/README.rst @@ -14,6 +14,26 @@ Install this extension using the following CLI command `az extension add --name Remove this extension using the following CLI command `az extension remove --name aks-preview`. +Azure Managed Lustre +-------------------- + +Install the ``k8s-extension`` CLI extension before enabling Azure Managed Lustre:: + + az extension add --name k8s-extension + az aks create --resource-group --name --generate-ssh-keys --enable-azure-managed-lustre + +For an existing AKS cluster, enable or disable the extension with:: + + az aks update --resource-group --name --enable-azure-managed-lustre + az aks update --resource-group --name --disable-azure-managed-lustre + +Enabling installs ``microsoft.azuremanagedlustre`` as ``azurelustre`` with cluster scope, +version ``0.6.0``, release train ``stable``, and automatic upgrades disabled. +As with Azure Container Storage, these operations wait for cluster provisioning before +installing or uninstalling the extension, even when ``--no-wait`` is specified. +The enable and disable flags cannot be used together. Disabling uninstalls the cluster +extension; it does not delete an Azure Managed Lustre file system. + Dependency between aks-preview and azure-cli/acs (azure-cli-core) ================================================================= diff --git a/src/aks-preview/azext_aks_preview/_params.py b/src/aks-preview/azext_aks_preview/_params.py index 8409f18d526..6ca6a028d98 100644 --- a/src/aks-preview/azext_aks_preview/_params.py +++ b/src/aks-preview/azext_aks_preview/_params.py @@ -1289,6 +1289,13 @@ def load_arguments(self, _): is_preview=True, action="store_true" ) + c.argument( + "enable_azure_managed_lustre", + action="store_true", + is_preview=True, + help="Install the Azure Managed Lustre extension version 0.6.0 on the stable release train " + "with automatic upgrades disabled. Requires the k8s-extension CLI extension.", + ) # azure container storage c.argument( "enable_azure_container_storage", @@ -2037,6 +2044,19 @@ def load_arguments(self, _): c.argument("disable_cost_analysis", action="store_true") c.argument('enable_ai_toolchain_operator', is_preview=True, action='store_true') c.argument('disable_ai_toolchain_operator', is_preview=True, action='store_true') + c.argument( + "enable_azure_managed_lustre", + action="store_true", + is_preview=True, + help="Install the Azure Managed Lustre extension version 0.6.0 on the stable release train " + "with automatic upgrades disabled. Requires the k8s-extension CLI extension.", + ) + c.argument( + "disable_azure_managed_lustre", + action="store_true", + is_preview=True, + help="Uninstall the Azure Managed Lustre extension from the cluster.", + ) # azure container storage c.argument( "enable_azure_container_storage", diff --git a/src/aks-preview/azext_aks_preview/azuremanagedlustre/__init__.py b/src/aks-preview/azext_aks_preview/azuremanagedlustre/__init__.py new file mode 100644 index 00000000000..34913fb394d --- /dev/null +++ b/src/aks-preview/azext_aks_preview/azuremanagedlustre/__init__.py @@ -0,0 +1,4 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- diff --git a/src/aks-preview/azext_aks_preview/azuremanagedlustre/_consts.py b/src/aks-preview/azext_aks_preview/azuremanagedlustre/_consts.py new file mode 100644 index 00000000000..3654da6d72a --- /dev/null +++ b/src/aks-preview/azext_aks_preview/azuremanagedlustre/_consts.py @@ -0,0 +1,9 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +CONST_AML_EXT_INSTALLATION_NAME = "azurelustre" +CONST_AML_K8S_EXTENSION_NAME = "microsoft.azuremanagedlustre" +CONST_AML_VERSION = "0.6.0" +CONST_AML_RELEASE_TRAIN = "stable" diff --git a/src/aks-preview/azext_aks_preview/azuremanagedlustre/_helpers.py b/src/aks-preview/azext_aks_preview/azuremanagedlustre/_helpers.py new file mode 100644 index 00000000000..f6ef1cf63e7 --- /dev/null +++ b/src/aks-preview/azext_aks_preview/azuremanagedlustre/_helpers.py @@ -0,0 +1,38 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +from azure.cli.core.azclierror import InvalidArgumentValueError, ResourceNotFoundError +from azext_aks_preview._consts import ( + CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME, + CONST_K8S_EXTENSION_CUSTOM_MOD_NAME, +) +from azext_aks_preview._helpers import get_k8s_extension_module +from azext_aks_preview.azuremanagedlustre._consts import ( + CONST_AML_EXT_INSTALLATION_NAME, + CONST_AML_K8S_EXTENSION_NAME, +) + + +def get_azure_managed_lustre_extension_client(cmd): + client_factory = get_k8s_extension_module(CONST_K8S_EXTENSION_CLIENT_FACTORY_MOD_NAME) + custom_module = get_k8s_extension_module(CONST_K8S_EXTENSION_CUSTOM_MOD_NAME) + return client_factory.cf_k8s_extension_operation(cmd.cli_ctx), custom_module + + +def check_if_extension_is_installed(cmd, resource_group, cluster_name): + client, custom_module = get_azure_managed_lustre_extension_client(cmd) + try: + extension = custom_module.show_k8s_extension( + client, resource_group, cluster_name, CONST_AML_EXT_INSTALLATION_NAME, "managedClusters" + ) + except ResourceNotFoundError: + return False + + if extension.extension_type.lower() != CONST_AML_K8S_EXTENSION_NAME: + raise InvalidArgumentValueError( + f"The extension '{CONST_AML_EXT_INSTALLATION_NAME}' already exists with type " + f"'{extension.extension_type}', not '{CONST_AML_K8S_EXTENSION_NAME}'." + ) + return True diff --git a/src/aks-preview/azext_aks_preview/azuremanagedlustre/_validators.py b/src/aks-preview/azext_aks_preview/azuremanagedlustre/_validators.py new file mode 100644 index 00000000000..ab3f3d8eec4 --- /dev/null +++ b/src/aks-preview/azext_aks_preview/azuremanagedlustre/_validators.py @@ -0,0 +1,17 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +from azure.cli.core.azclierror import InvalidArgumentValueError, MutuallyExclusiveArgumentError + + +def validate_azure_managed_lustre_params(enable, disable, is_extension_installed=None): + if enable and disable: + raise MutuallyExclusiveArgumentError( + "Cannot set --enable-azure-managed-lustre and --disable-azure-managed-lustre together." + ) + if disable and is_extension_installed is False: + raise InvalidArgumentValueError( + "Cannot set --disable-azure-managed-lustre. Azure Managed Lustre is not enabled in the cluster." + ) diff --git a/src/aks-preview/azext_aks_preview/azuremanagedlustre/aml_ops.py b/src/aks-preview/azext_aks_preview/azuremanagedlustre/aml_ops.py new file mode 100644 index 00000000000..a98530d34f9 --- /dev/null +++ b/src/aks-preview/azext_aks_preview/azuremanagedlustre/aml_ops.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. +# -------------------------------------------------------------------------------------------- + +from azure.cli.core.azclierror import UnknownError +from azure.cli.core.commands import LongRunningOperation +from knack.log import get_logger + +from azext_aks_preview.azuremanagedlustre._consts import ( + CONST_AML_EXT_INSTALLATION_NAME, + CONST_AML_K8S_EXTENSION_NAME, + CONST_AML_RELEASE_TRAIN, + CONST_AML_VERSION, +) +from azext_aks_preview.azuremanagedlustre._helpers import get_azure_managed_lustre_extension_client + +logger = get_logger(__name__) + + +def perform_enable_azure_managed_lustre(cmd, resource_group, cluster_name): + client, custom_module = get_azure_managed_lustre_extension_client(cmd) + result = custom_module.create_k8s_extension( + cmd, + client, + resource_group, + cluster_name, + CONST_AML_EXT_INSTALLATION_NAME, + "managedClusters", + CONST_AML_K8S_EXTENSION_NAME, + version=CONST_AML_VERSION, + release_train=CONST_AML_RELEASE_TRAIN, + scope="cluster", + auto_upgrade_minor_version=False, + ) + extension = LongRunningOperation(cmd.cli_ctx)(result) + if extension is None or extension.provisioning_state != "Succeeded": + raise UnknownError("Azure Managed Lustre extension installation did not succeed.") + logger.warning("Azure Managed Lustre successfully installed.") + + +def perform_disable_azure_managed_lustre(cmd, resource_group, cluster_name): + client, custom_module = get_azure_managed_lustre_extension_client(cmd) + result = custom_module.delete_k8s_extension( + cmd, + client, + resource_group, + cluster_name, + CONST_AML_EXT_INSTALLATION_NAME, + "managedClusters", + yes=True, + ) + # The delegated delete returns None when it cannot retrieve the extension. + if result is None: + raise UnknownError("Azure Managed Lustre extension could not be retrieved for deletion.") + LongRunningOperation(cmd.cli_ctx)(result) + logger.warning("Azure Managed Lustre has been disabled.") diff --git a/src/aks-preview/azext_aks_preview/custom.py b/src/aks-preview/azext_aks_preview/custom.py index 187811c0345..015d7478c8d 100644 --- a/src/aks-preview/azext_aks_preview/custom.py +++ b/src/aks-preview/azext_aks_preview/custom.py @@ -1440,6 +1440,7 @@ def aks_create( enable_cost_analysis=False, # AI toolchain operator enable_ai_toolchain_operator=False, + enable_azure_managed_lustre=False, # azure container storage enable_azure_container_storage=None, container_storage_version=None, @@ -1714,6 +1715,8 @@ def aks_update( # AI toolchain operator enable_ai_toolchain_operator=False, disable_ai_toolchain_operator=False, + enable_azure_managed_lustre=False, + disable_azure_managed_lustre=False, # azure container storage enable_azure_container_storage=None, disable_azure_container_storage=None, diff --git a/src/aks-preview/azext_aks_preview/managed_cluster_decorator.py b/src/aks-preview/azext_aks_preview/managed_cluster_decorator.py index a9a88b3c8e8..12a777f194e 100644 --- a/src/aks-preview/azext_aks_preview/managed_cluster_decorator.py +++ b/src/aks-preview/azext_aks_preview/managed_cluster_decorator.py @@ -96,6 +96,15 @@ perform_enable_azure_container_storage_v1, perform_azure_container_storage_update, ) +from azext_aks_preview.azuremanagedlustre._helpers import ( + check_if_extension_is_installed as check_if_azure_managed_lustre_is_installed, + get_azure_managed_lustre_extension_client, +) +from azext_aks_preview.azuremanagedlustre._validators import validate_azure_managed_lustre_params +from azext_aks_preview.azuremanagedlustre.aml_ops import ( + perform_disable_azure_managed_lustre, + perform_enable_azure_managed_lustre, +) from azext_aks_preview.azuremonitormetrics.azuremonitorprofile import ( ensure_azure_monitor_profile_prerequisites, ) @@ -291,6 +300,8 @@ def external_functions(self) -> SimpleNamespace: "perform_disable_azure_container_storage_v1" ] = perform_disable_azure_container_storage_v1 external_functions["perform_azure_container_storage_update"] = perform_azure_container_storage_update + external_functions["perform_enable_azure_managed_lustre"] = perform_enable_azure_managed_lustre + external_functions["perform_disable_azure_managed_lustre"] = perform_disable_azure_managed_lustre external_functions["sanitize_loganalytics_ws_resource_id"] = sanitize_loganalytics_ws_resource_id # Override base module function with preview version that uses REST API to avoid # "Request Header Fields Too Large" errors @@ -5358,6 +5369,13 @@ def set_up_azure_monitor_profile(self, mc: ManagedCluster) -> ManagedCluster: return mc + def set_up_azure_managed_lustre(self, mc: ManagedCluster) -> ManagedCluster: + self._ensure_mc(mc) + if self.context.raw_param.get("enable_azure_managed_lustre"): + get_azure_managed_lustre_extension_client(self.cmd) + self.context.set_intermediate("enable_azure_managed_lustre", True, overwrite_exists=True) + return mc + def set_up_azure_container_storage(self, mc: ManagedCluster) -> ManagedCluster: # pylint: disable=too-many-locals """Set up azure container storage for the Managed Cluster object :return: ManagedCluster @@ -5952,6 +5970,7 @@ def construct_mc_profile_preview(self, bypass_restore_defaults: bool = False) -> mc = self.set_up_ai_toolchain_operator(mc) # set up for azure container storage mc = self.set_up_azure_container_storage(mc) + mc = self.set_up_azure_managed_lustre(mc) # set up node provisioning profile mc = self.set_up_node_provisioning_profile(mc) # set up node disruption policy @@ -6044,6 +6063,7 @@ def check_is_postprocessing_required(self, mc: ManagedCluster) -> bool: (enable_managed_identity and attach_acr) or need_grant_vnet_permission_to_cluster_identity or enable_azure_container_storage or + self.context.get_intermediate("enable_azure_managed_lustre", default_value=False) or enable_backup ): return True @@ -6303,6 +6323,11 @@ def postprocessing_after_mc_created(self, cluster: ManagedCluster) -> None: is_called_from_extension=True, ) + if self.context.get_intermediate("enable_azure_managed_lustre", default_value=False): + self.context.external_functions.perform_enable_azure_managed_lustre( + self.cmd, self.context.get_resource_group_name(), self.context.get_name() + ) + # Add role assignments for automatic sku if cluster.sku is not None and cluster.sku.name == "Automatic": try: @@ -6733,6 +6758,20 @@ def update_monitoring_profile_flow_logs(self, mc: ManagedCluster) -> ManagedClus return mc + def update_azure_managed_lustre(self, mc: ManagedCluster) -> ManagedCluster: + self._ensure_mc(mc) + enable = self.context.raw_param.get("enable_azure_managed_lustre", False) + disable = self.context.raw_param.get("disable_azure_managed_lustre", False) + validate_azure_managed_lustre_params(enable, disable) + if enable or disable: + installed = check_if_azure_managed_lustre_is_installed( + self.cmd, self.context.get_resource_group_name(), self.context.get_name() + ) + validate_azure_managed_lustre_params(enable, disable, installed) + self.context.set_intermediate("enable_azure_managed_lustre", enable, overwrite_exists=True) + self.context.set_intermediate("disable_azure_managed_lustre", disable, overwrite_exists=True) + return mc + # pylint: disable=too-many-statements,too-many-locals,too-many-branches def update_azure_container_storage(self, mc: ManagedCluster) -> ManagedCluster: """Update azure container storage for the Managed Cluster object @@ -9106,6 +9145,7 @@ def update_mc_profile_preview(self) -> ManagedCluster: mc = self.update_ai_toolchain_operator(mc) # update azure container storage mc = self.update_azure_container_storage(mc) + mc = self.update_azure_managed_lustre(mc) # update node provisioning profile mc = self.update_node_provisioning_profile(mc) # update bootstrap profile @@ -9157,6 +9197,8 @@ def check_is_postprocessing_required(self, mc: ManagedCluster) -> bool: # Note: monitoring_addon_disable_postprocessing_required is no longer used - cleanup is done upfront # pylint: disable=too-many-boolean-expressions if (enable_azure_container_storage or disable_azure_container_storage) or \ + self.context.get_intermediate("enable_azure_managed_lustre", default_value=False) or \ + self.context.get_intermediate("disable_azure_managed_lustre", default_value=False) or \ (keyvault_id and enable_azure_keyvault_secrets_provider_addon) or \ (monitoring_addon_postprocessing_required) or \ enable_backup: @@ -9172,6 +9214,15 @@ def postprocessing_after_mc_created(self, cluster: ManagedCluster) -> None: """ super().postprocessing_after_mc_created(cluster) + if self.context.get_intermediate("enable_azure_managed_lustre", default_value=False): + self.context.external_functions.perform_enable_azure_managed_lustre( + self.cmd, self.context.get_resource_group_name(), self.context.get_name() + ) + if self.context.get_intermediate("disable_azure_managed_lustre", default_value=False): + self.context.external_functions.perform_disable_azure_managed_lustre( + self.cmd, self.context.get_resource_group_name(), self.context.get_name() + ) + # Handle monitoring addon postprocessing (enable case) monitoring_addon_postprocessing_required = self.context.get_intermediate( "monitoring_addon_postprocessing_required", default_value=False diff --git a/src/aks-preview/azext_aks_preview/tests/latest/test_azure_managed_lustre.py b/src/aks-preview/azext_aks_preview/tests/latest/test_azure_managed_lustre.py new file mode 100644 index 00000000000..cf597cc72e8 --- /dev/null +++ b/src/aks-preview/azext_aks_preview/tests/latest/test_azure_managed_lustre.py @@ -0,0 +1,339 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import inspect +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from azure.core.exceptions import HttpResponseError +from azure.cli.core.parser import AzCliCommandParser +from azure.cli.core.azclierror import ( + InvalidArgumentValueError, + MutuallyExclusiveArgumentError, + ResourceNotFoundError, + UnknownError, +) +from azext_aks_preview import ContainerServiceCommandsLoader, register_aks_preview_resource_type +from azext_aks_preview._client_factory import CUSTOM_MGMT_AKS_PREVIEW +from azext_aks_preview.azuremanagedlustre._helpers import ( + check_if_extension_is_installed, + get_azure_managed_lustre_extension_client, +) +from azext_aks_preview.azuremanagedlustre._validators import validate_azure_managed_lustre_params +from azext_aks_preview.azuremanagedlustre.aml_ops import ( + perform_disable_azure_managed_lustre, + perform_enable_azure_managed_lustre, +) +from azext_aks_preview.custom import aks_create, aks_update +from azext_aks_preview.managed_cluster_decorator import ( + AKSPreviewManagedClusterCreateDecorator, + AKSPreviewManagedClusterModels, + AKSPreviewManagedClusterUpdateDecorator, +) +from azext_aks_preview.tests.latest.mocks import MockCLI, MockCmd + + +HELPERS = "azext_aks_preview.azuremanagedlustre._helpers" +OPS = "azext_aks_preview.azuremanagedlustre.aml_ops" +DECORATOR = "azext_aks_preview.managed_cluster_decorator" + + +class AzureManagedLustreHelpersTestCase(unittest.TestCase): + def setUp(self): + self.cmd = Mock() + self.client = Mock() + self.custom = Mock() + self.client_patch = patch( + HELPERS + ".get_azure_managed_lustre_extension_client", + return_value=(self.client, self.custom), + ) + + def test_get_extension_client_reuses_module_loader(self): + factory = Mock() + with patch(HELPERS + ".get_k8s_extension_module", side_effect=[factory, self.custom]) as loader: + client, custom = get_azure_managed_lustre_extension_client(self.cmd) + self.assertEqual([call.args[0] for call in loader.call_args_list], [ + "azext_k8s_extension._client_factory", "azext_k8s_extension.custom", + ]) + factory.cf_k8s_extension_operation.assert_called_once_with(self.cmd.cli_ctx) + self.assertIs(client, factory.cf_k8s_extension_operation.return_value) + self.assertIs(custom, self.custom) + + def test_missing_cli_extension_is_reported(self): + with patch(HELPERS + ".get_k8s_extension_module", side_effect=UnknownError("Install k8s-extension")): + with self.assertRaisesRegex(UnknownError, "Install k8s-extension"): + get_azure_managed_lustre_extension_client(self.cmd) + + def test_installed_extension_type_is_case_insensitive(self): + self.custom.show_k8s_extension.return_value.extension_type = "Microsoft.AzureManagedLustre" + with self.client_patch: + self.assertTrue(check_if_extension_is_installed(self.cmd, "rg", "cluster")) + self.custom.show_k8s_extension.assert_called_once_with( + self.client, "rg", "cluster", "azurelustre", "managedClusters" + ) + + def test_missing_extension(self): + self.custom.show_k8s_extension.side_effect = ResourceNotFoundError("not found") + with self.client_patch: + self.assertFalse(check_if_extension_is_installed(self.cmd, "rg", "cluster")) + + def test_name_collision_is_rejected(self): + self.custom.show_k8s_extension.return_value.extension_type = "microsoft.other" + with self.client_patch, self.assertRaisesRegex(InvalidArgumentValueError, "microsoft.other"): + check_if_extension_is_installed(self.cmd, "rg", "cluster") + + def test_service_errors_are_not_treated_as_missing(self): + self.custom.show_k8s_extension.side_effect = HttpResponseError("Forbidden") + with self.client_patch, self.assertRaisesRegex(HttpResponseError, "Forbidden"): + check_if_extension_is_installed(self.cmd, "rg", "cluster") + + +class AzureManagedLustreValidatorsTestCase(unittest.TestCase): + def test_conflicting_flags(self): + with self.assertRaises(MutuallyExclusiveArgumentError): + validate_azure_managed_lustre_params(True, True) + + def test_disable_requires_installed_extension(self): + with self.assertRaisesRegex(InvalidArgumentValueError, "not enabled"): + validate_azure_managed_lustre_params(False, True, False) + + def test_valid_parameters(self): + for enable, disable, installed in [ + (False, False, None), (True, False, False), (True, False, True), (False, True, True), + ]: + with self.subTest(enable=enable, disable=disable, installed=installed): + validate_azure_managed_lustre_params(enable, disable, installed) + + +class AzureManagedLustreOperationsTestCase(unittest.TestCase): + def setUp(self): + self.cmd = Mock() + self.client = Mock() + self.custom = Mock() + self.client_patch = patch( + OPS + ".get_azure_managed_lustre_extension_client", + return_value=(self.client, self.custom), + ) + self.client_patch.start() + self.addCleanup(self.client_patch.stop) + self.waiter = Mock(return_value=SimpleNamespace(provisioning_state="Succeeded")) + self.wait_patch = patch(OPS + ".LongRunningOperation", return_value=self.waiter) + self.wait_patch.start() + self.addCleanup(self.wait_patch.stop) + + def test_enable_matches_requested_installation(self): + perform_enable_azure_managed_lustre(self.cmd, "rg", "cluster") + self.custom.create_k8s_extension.assert_called_once_with( + self.cmd, self.client, "rg", "cluster", "azurelustre", "managedClusters", + "microsoft.azuremanagedlustre", version="0.6.0", release_train="stable", + scope="cluster", auto_upgrade_minor_version=False, + ) + self.waiter.assert_called_once_with(self.custom.create_k8s_extension.return_value) + + def test_enable_propagates_service_error_without_deleting_extension(self): + self.custom.create_k8s_extension.side_effect = HttpResponseError("Forbidden") + with self.assertRaisesRegex(HttpResponseError, "Forbidden"): + perform_enable_azure_managed_lustre(self.cmd, "rg", "cluster") + self.custom.delete_k8s_extension.assert_not_called() + + def test_enable_reports_unsuccessful_provisioning(self): + for result in [None, SimpleNamespace(provisioning_state="Failed"), + SimpleNamespace(provisioning_state="Canceled")]: + with self.subTest(result=result): + self.waiter.return_value = result + with self.assertRaisesRegex(UnknownError, "did not succeed"): + perform_enable_azure_managed_lustre(self.cmd, "rg", "cluster") + + def test_enable_propagates_polling_error(self): + self.waiter.side_effect = HttpResponseError("Provisioning failed") + with self.assertRaisesRegex(HttpResponseError, "Provisioning failed"): + perform_enable_azure_managed_lustre(self.cmd, "rg", "cluster") + + def test_disable_waits_for_deletion(self): + perform_disable_azure_managed_lustre(self.cmd, "rg", "cluster") + self.custom.delete_k8s_extension.assert_called_once_with( + self.cmd, self.client, "rg", "cluster", "azurelustre", "managedClusters", yes=True + ) + self.waiter.assert_called_once_with(self.custom.delete_k8s_extension.return_value) + + def test_disable_reports_failed_retrieval(self): + self.custom.delete_k8s_extension.return_value = None + with self.assertRaisesRegex(UnknownError, "could not be retrieved"): + perform_disable_azure_managed_lustre(self.cmd, "rg", "cluster") + self.waiter.assert_not_called() + + def test_disable_propagates_polling_error(self): + self.waiter.side_effect = HttpResponseError("Deletion failed") + with self.assertRaisesRegex(HttpResponseError, "Deletion failed"): + perform_disable_azure_managed_lustre(self.cmd, "rg", "cluster") + + +class AzureManagedLustreDecoratorTestCase(unittest.TestCase): + def setUp(self): + register_aks_preview_resource_type() + self.cmd = MockCmd(MockCLI()) + self.models = AKSPreviewManagedClusterModels(self.cmd, CUSTOM_MGMT_AKS_PREVIEW) + self.client = Mock() + + def make_decorator(self, create, **params): + decorator_type = ( + AKSPreviewManagedClusterCreateDecorator if create else AKSPreviewManagedClusterUpdateDecorator + ) + raw_params = {"resource_group_name": "rg", "name": "cluster"} + raw_params.update(params) + decorator = decorator_type(self.cmd, self.client, raw_params, CUSTOM_MGMT_AKS_PREVIEW) + decorator.context.attach_mc(self.models.ManagedCluster(location="eastus")) + decorator.context.set_intermediate("subscription_id", "test-subscription") + return decorator + + def test_command_signatures_keep_feature_opt_in(self): + for command in [aks_create, aks_update]: + self.assertIs(inspect.signature(command).parameters["enable_azure_managed_lustre"].default, False) + self.assertIs(inspect.signature(aks_update).parameters["disable_azure_managed_lustre"].default, False) + self.assertNotIn("disable_azure_managed_lustre", inspect.signature(aks_create).parameters) + + def test_cli_arguments_are_registered(self): + self.cmd.cli_ctx.local_context = Mock(is_on=False) + for command_name in ["aks create", "aks update"]: + with self.subTest(command=command_name): + parser = AzCliCommandParser(cli_ctx=self.cmd.cli_ctx) + self.cmd.cli_ctx.invocation = Mock(parser=parser, data={"command_string": command_name}) + loader = ContainerServiceCommandsLoader(self.cmd.cli_ctx) + loader.load_command_table(command_name.split()) + command = loader.command_table[command_name] + command.load_arguments() + loader.load_arguments(command_name) + loader._apply_parameter_info(command_name, command) # pylint: disable=protected-access + loader.command_table = {command_name: command} + parser.load_command_table(loader) + args = command_name.split() + ["--resource-group", "rg", "--name", "cluster"] + enable = command.arguments["enable_azure_managed_lustre"] + self.assertEqual(enable.options_list, ["--enable-azure-managed-lustre"]) + self.assertFalse(parser.parse_args(args).enable_azure_managed_lustre) + self.assertTrue(parser.parse_args(args + enable.options_list).enable_azure_managed_lustre) + if command_name == "aks update": + disable = command.arguments["disable_azure_managed_lustre"] + self.assertEqual(disable.options_list, ["--disable-azure-managed-lustre"]) + self.assertFalse(parser.parse_args(args).disable_azure_managed_lustre) + self.assertTrue(parser.parse_args(args + disable.options_list).disable_azure_managed_lustre) + else: + self.assertNotIn("disable_azure_managed_lustre", command.arguments) + + def test_external_functions(self): + decorator = self.make_decorator(True) + functions = decorator.context.external_functions + self.assertIs(functions.perform_enable_azure_managed_lustre, perform_enable_azure_managed_lustre) + self.assertIs(functions.perform_disable_azure_managed_lustre, perform_disable_azure_managed_lustre) + + def test_create_preflights_dependency_without_installing(self): + decorator = self.make_decorator(True, enable_azure_managed_lustre=True) + mc = decorator.context.mc + with patch(DECORATOR + ".get_azure_managed_lustre_extension_client") as dependency: + self.assertIs(decorator.set_up_azure_managed_lustre(mc), mc) + dependency.assert_called_once_with(self.cmd) + self.assertTrue(decorator.context.get_intermediate("enable_azure_managed_lustre")) + self.assertTrue(decorator.check_is_postprocessing_required(mc)) + + def test_create_missing_dependency_fails_before_cluster_request(self): + decorator = self.make_decorator(True, enable_azure_managed_lustre=True) + with patch(DECORATOR + ".get_azure_managed_lustre_extension_client", + side_effect=UnknownError("Install k8s-extension")): + with self.assertRaisesRegex(UnknownError, "Install k8s-extension"): + decorator.set_up_azure_managed_lustre(decorator.context.mc) + self.client.begin_create_or_update.assert_not_called() + + def test_omitted_flags_do_not_load_extension_or_request_postprocessing(self): + for create in [True, False]: + with self.subTest(create=create): + decorator = self.make_decorator(create) + mc = decorator.context.mc + with patch(DECORATOR + ".get_azure_managed_lustre_extension_client") as dependency, \ + patch(DECORATOR + ".check_if_azure_managed_lustre_is_installed") as installed: + setup = decorator.set_up_azure_managed_lustre if create else decorator.update_azure_managed_lustre + self.assertIs(setup(mc), mc) + dependency.assert_not_called() + installed.assert_not_called() + self.assertFalse(decorator.check_is_postprocessing_required(mc)) + + def test_update_rejects_conflicting_flags_before_lookup(self): + decorator = self.make_decorator( + False, enable_azure_managed_lustre=True, disable_azure_managed_lustre=True + ) + with patch(DECORATOR + ".check_if_azure_managed_lustre_is_installed") as installed: + with self.assertRaises(MutuallyExclusiveArgumentError): + decorator.update_azure_managed_lustre(decorator.context.mc) + installed.assert_not_called() + + def test_update_rejects_disable_when_not_installed(self): + decorator = self.make_decorator(False, disable_azure_managed_lustre=True) + with patch(DECORATOR + ".check_if_azure_managed_lustre_is_installed", return_value=False): + with self.assertRaisesRegex(InvalidArgumentValueError, "not enabled"): + decorator.update_azure_managed_lustre(decorator.context.mc) + + def test_update_propagates_lookup_errors_before_cluster_request(self): + decorator = self.make_decorator(False, enable_azure_managed_lustre=True) + with patch(DECORATOR + ".check_if_azure_managed_lustre_is_installed", + side_effect=HttpResponseError("Forbidden")): + with self.assertRaisesRegex(HttpResponseError, "Forbidden"): + decorator.update_azure_managed_lustre(decorator.context.mc) + self.client.begin_create_or_update.assert_not_called() + + def test_update_prepares_enable_and_disable(self): + for enable, installed in [(True, False), (True, True), (False, True)]: + with self.subTest(enable=enable, installed=installed): + decorator = self.make_decorator( + False, enable_azure_managed_lustre=enable, disable_azure_managed_lustre=not enable + ) + mc = decorator.context.mc + with patch(DECORATOR + ".check_if_azure_managed_lustre_is_installed", + return_value=installed) as lookup: + self.assertIs(decorator.update_azure_managed_lustre(mc), mc) + lookup.assert_called_once_with(self.cmd, "rg", "cluster") + self.assertEqual(decorator.context.get_intermediate("enable_azure_managed_lustre"), enable) + self.assertEqual(decorator.context.get_intermediate("disable_azure_managed_lustre"), not enable) + self.assertTrue(decorator.check_is_postprocessing_required(mc)) + + def test_postprocessing_performs_requested_operation(self): + for create, enable in [(True, True), (False, True), (False, False)]: + with self.subTest(create=create, enable=enable): + decorator = self.make_decorator(create) + operation = "enable" if enable else "disable" + decorator.context.set_intermediate(operation + "_azure_managed_lustre", True) + functions = decorator.context.external_functions + with patch.object(functions, "perform_enable_azure_managed_lustre") as install, \ + patch.object(functions, "perform_disable_azure_managed_lustre") as uninstall, \ + patch("azure.cli.command_modules.acs.managed_cluster_decorator." + "AKSManagedClusterUpdateDecorator.postprocessing_after_mc_created"): + decorator.postprocessing_after_mc_created(decorator.context.mc) + (install if enable else uninstall).assert_called_once_with(self.cmd, "rg", "cluster") + (uninstall if enable else install).assert_not_called() + + def test_put_waits_for_cluster_before_lustre_even_with_no_wait(self): + for create in [True, False]: + with self.subTest(create=create): + decorator = self.make_decorator(create, no_wait=True) + decorator.context.set_intermediate("enable_azure_managed_lustre", True) + mc = decorator.context.mc + events = [] + cluster = self.models.ManagedCluster(location="eastus") + + def wait_for_cluster(poller): + self.assertIs(poller, self.client.begin_create_or_update.return_value) + events.append("cluster-ready") + return cluster + + def install(cmd, resource_group, cluster_name): + self.assertEqual(events, ["cluster-ready"]) + events.append("extension-installed") + + with patch(DECORATOR + ".LongRunningOperation", return_value=wait_for_cluster), \ + patch.object(decorator, "immediate_processing_after_request"), \ + patch.object(decorator.context.external_functions, "perform_enable_azure_managed_lustre", + side_effect=install), \ + patch("azure.cli.command_modules.acs.managed_cluster_decorator." + "AKSManagedClusterUpdateDecorator.postprocessing_after_mc_created"): + self.assertIs(decorator.put_mc(mc), cluster) + self.assertEqual(events, ["cluster-ready", "extension-installed"])