Adding local dns profile to az aks nodepool update

src/aks-preview/azext_aks_preview/_consts.py

@@ -372,3 +372,8 @@
 CONST_K8S_EXTENSION_NAME = "k8s-extension"
 CONST_K8S_EXTENSION_ACTION_MOD_NAME = "azext_k8s_extension.action"
 CONST_K8S_EXTENSION_FORMAT_MOD_NAME = "azext_k8s_extension._format"
+
+# Local DNS profile modes
+CONST_LOCAL_DNS_MODE_REQUIRED = "required"
+CONST_LOCAL_DNS_MODE_PREFERRED = "preferred"
+CONST_LOCAL_DNS_MODE_DISABLED = "disabled"

src/aks-preview/azext_aks_preview/_params.py

@@ -65,6 +65,9 @@
     CONST_GPU_INSTANCE_PROFILE_MIG7_G,
     CONST_LOAD_BALANCER_SKU_BASIC,
     CONST_LOAD_BALANCER_SKU_STANDARD,
+    CONST_LOCAL_DNS_MODE_REQUIRED,
+    CONST_LOCAL_DNS_MODE_PREFERRED,
+    CONST_LOCAL_DNS_MODE_DISABLED,
     CONST_MANAGED_CLUSTER_SKU_TIER_FREE,
     CONST_MANAGED_CLUSTER_SKU_TIER_STANDARD,
     CONST_MANAGED_CLUSTER_SKU_TIER_PREMIUM,
@@ -1823,6 +1826,29 @@ def load_arguments(self, _):
             arg_type=get_enum_type(node_os_skus_update),
             validator=validate_os_sku,
         )
+        # local dns
+        c.argument(
+            "set_localdns",
+            options_list=["--set-localdns"],
+            action="store_true",
+            help="Enable or update the local DNS profile for this nodepool. Must be used with --localdns-mode or --localdns-config."
+        )
+        c.argument(
+            "localdns_config",
+            options_list=["--localdns-config"],
+            type=file_type,
+            help="Path to a JSON file with the local DNS profile configuration. Must be used with --set-localdns."
+        )
+        c.argument(
+            "localdns_mode",
+            options_list=["--localdns-mode"],
+            arg_type=get_enum_type([
+                CONST_LOCAL_DNS_MODE_REQUIRED,
+                CONST_LOCAL_DNS_MODE_PREFERRED,
+                CONST_LOCAL_DNS_MODE_DISABLED
+            ]),
+            help="The mode for the local DNS profile. Must be used with --set-localdns."
+        )
         # In update scenario, use emtpy str as default.
         c.argument('ssh_access', arg_type=get_enum_type(ssh_accesses), is_preview=True)
         c.argument('yes', options_list=['--yes', '-y'], help='Do not prompt for confirmation.', action='store_true')

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -41,7 +41,12 @@
     CONST_DEFAULT_WINDOWS_VMS_VM_SIZE,
     CONST_SSH_ACCESS_LOCALUSER,
     CONST_GPU_DRIVER_NONE,
+    CONST_LOCAL_DNS_MODE_REQUIRED,
+    CONST_LOCAL_DNS_MODE_PREFERRED,
+    CONST_LOCAL_DNS_MODE_DISABLED,
 )
+import json
+
 from azext_aks_preview._helpers import (
     get_nodepool_snapshot_by_snapshot_id,
     filter_hard_taints,
@@ -808,6 +813,49 @@ def get_disable_fips_image(self) -> bool:
         # read the original value passed by the command
         return self.raw_param.get("disable_fips_image")
 
+    def get_set_localdns(self) -> bool:
+        return self.raw_param.get("set_localdns", False)
+
+    def get_localdns_config(self):
+        return self.raw_param.get("localdns_config")
+
+    def get_localdns_mode(self):
+        return self.raw_param.get("localdns_mode")
+
+    def get_localdns_profile(self):
+        """
+        Returns the local DNS profile dict if set, or None.
+        Handles validation and file loading.
+        """
+        set_localdns = self.get_set_localdns()
+        config = self.get_localdns_config()
+        mode = self.get_localdns_mode()
+        if not set_localdns and (config or mode):
+            raise InvalidArgumentValueError("--set-localdns must be specified when using --localdns-config or --localdns-mode.")
+        if set_localdns:
+            if not (config or mode):
+                raise InvalidArgumentValueError("--set-localdns requires either --localdns-config or --localdns-mode.")
+            if config and mode:
+                raise MutuallyExclusiveArgumentError("Cannot specify both --localdns-config and --localdns-mode.")
+            if config:
+                # Load and validate JSON file
+                try:
+                    if isinstance(config, str):
+                        with open(config, "r") as f:
+                            profile = json.load(f)
+                    else:
+                        # If already file-like (from CLI file_type)
+                        profile = json.load(config)
+                except Exception as ex:
+                    raise InvalidArgumentValueError(f"Failed to load local DNS config file: {ex}")
+                return profile
+            if mode:
+                if mode not in [CONST_LOCAL_DNS_MODE_REQUIRED, CONST_LOCAL_DNS_MODE_PREFERRED, CONST_LOCAL_DNS_MODE_DISABLED]:
+                    raise InvalidArgumentValueError(f"Invalid local DNS mode: {mode}")
+                # Return a minimal profile dict with just the mode
+                return {"mode": mode}
+        return None
+
 
 class AKSPreviewAgentPoolAddDecorator(AKSAgentPoolAddDecorator):
     def __init__(
@@ -1342,6 +1390,9 @@ def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -
         # update ssh access
         agentpool = self.update_ssh_access(agentpool)
 
+        # update locald DNS profile
+        localdns_profile = self.context.get_localdns_profile()
+
         return agentpool
 
     def update_upgrade_settings(self, agentpool: AgentPool) -> AgentPool:

src/aks-preview/azext_aks_preview/tests/latest/test_localdns_profile.py

@@ -0,0 +1,101 @@
+# --------------------------------------------------------------------------------------------
+# 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 unittest.mock import Mock
+from azext_aks_preview.agentpool_decorator import AKSPreviewAgentPoolContext, AKSPreviewAgentPoolModels, AKSPreviewAgentPoolUpdateDecorator
+from azure.cli.command_modules.acs.agentpool_decorator import AKSAgentPoolParamDict
+from azure.cli.command_modules.acs._consts import AgentPoolDecoratorMode, DecoratorMode
+from azure.cli.core.azclierror import InvalidArgumentValueError, MutuallyExclusiveArgumentError
+from types import SimpleNamespace
+import tempfile
+import json
+
+class TestLocalDNSProfile(unittest.TestCase):
+    def setUp(self):
+        self.cmd = Mock()
+        self.models = Mock()
+        self.models.AgentPoolLocalDNSProfile = lambda **kwargs: SimpleNamespace(**kwargs)
+        self.agentpool_decorator_mode = AgentPoolDecoratorMode.STANDALONE
+
+    def test_localdns_mode(self):
+        ctx = AKSPreviewAgentPoolContext(
+            self.cmd,
+            AKSAgentPoolParamDict({
+                "set_localdns": True,
+                "localdns_mode": "preferred",
+                "localdns_config": None
+            }),
+            self.models,
+            DecoratorMode.UPDATE,
+            self.agentpool_decorator_mode,
+        )
+        profile = ctx.get_localdns_profile()
+        self.assertEqual(profile, {"mode": "preferred"})
+
+    def test_localdns_config(self):
+        config = {"mode": "required", "custom": "foo"}
+        with tempfile.NamedTemporaryFile(mode="w+", delete=False) as f:
+            json.dump(config, f)
+            f.flush()
+            ctx = AKSPreviewAgentPoolContext(
+                self.cmd,
+                AKSAgentPoolParamDict({
+                    "set_localdns": True,
+                    "localdns_mode": None,
+                    "localdns_config": f.name
+                }),
+                self.models,
+                DecoratorMode.UPDATE,
+                self.agentpool_decorator_mode,
+            )
+            profile = ctx.get_localdns_profile()
+            self.assertEqual(profile, config)
+
+    def test_localdns_validation(self):
+        # Missing --set-localdns
+        ctx = AKSPreviewAgentPoolContext(
+            self.cmd,
+            AKSAgentPoolParamDict({
+                "set_localdns": False,
+                "localdns_mode": "preferred",
+                "localdns_config": None
+            }),
+            self.models,
+            DecoratorMode.UPDATE,
+            self.agentpool_decorator_mode,
+        )
+        with self.assertRaises(InvalidArgumentValueError):
+            ctx.get_localdns_profile()
+        # Both config and mode
+        ctx = AKSPreviewAgentPoolContext(
+            self.cmd,
+            AKSAgentPoolParamDict({
+                "set_localdns": True,
+                "localdns_mode": "preferred",
+                "localdns_config": "foo.json"
+            }),
+            self.models,
+            DecoratorMode.UPDATE,
+            self.agentpool_decorator_mode,
+        )
+        with self.assertRaises(MutuallyExclusiveArgumentError):
+            ctx.get_localdns_profile()
+        # Neither config nor mode
+        ctx = AKSPreviewAgentPoolContext(
+            self.cmd,
+            AKSAgentPoolParamDict({
+                "set_localdns": True,
+                "localdns_mode": None,
+                "localdns_config": None
+            }),
+            self.models,
+            DecoratorMode.UPDATE,
+            self.agentpool_decorator_mode,
+        )
+        with self.assertRaises(InvalidArgumentValueError):
+            ctx.get_localdns_profile()
+
+if __name__ == "__main__":
+    unittest.main()