diff --git a/src/azure-cli/azure/cli/command_modules/acs/_help.py b/src/azure-cli/azure/cli/command_modules/acs/_help.py index e0b7ae94fcd..3b34e767a5e 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/_help.py +++ b/src/azure-cli/azure/cli/command_modules/acs/_help.py @@ -2065,6 +2065,9 @@ - name: --asg-ids type: string short-summary: The IDs of the application security groups to which the node pool's network interface should belong. When specified, format should be a space-separated list of IDs. + - name: --enable-managed-dranet + type: bool + short-summary: Enable Managed DRANET on the node pool. - name: --node-public-ip-tags type: string short-summary: The ipTags of the node public IPs. @@ -2240,6 +2243,9 @@ - name: --asg-ids type: string short-summary: The IDs of the application security groups to which the node pool's network interface should belong. When specified, format should be a space-separated list of IDs. + - name: --enable-managed-dranet + type: bool + short-summary: Enable Managed DRANET on the node pool. - name: --os-sku type: string short-summary: The os-sku of the agent node pool. diff --git a/src/azure-cli/azure/cli/command_modules/acs/_params.py b/src/azure-cli/azure/cli/command_modules/acs/_params.py index 09056477334..c73a1e0c872 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/_params.py +++ b/src/azure-cli/azure/cli/command_modules/acs/_params.py @@ -1160,6 +1160,7 @@ def load_arguments(self, _): c.argument('gpu_instance_profile', arg_type=get_enum_type(gpu_instance_profiles)) c.argument('allowed_host_ports', nargs='+', validator=validate_allowed_host_ports) c.argument('asg_ids', nargs='+', validator=validate_application_security_groups) + c.argument('enable_managed_dranet', action='store_true') c.argument('node_public_ip_tags', arg_type=tags_type, validator=validate_node_public_ip_tags, help='space-separated tags: key[=value] [key[=value] ...].') c.argument("message_of_the_day", validator=validate_message_of_the_day) @@ -1194,6 +1195,7 @@ def load_arguments(self, _): c.argument('scale_down_mode', arg_type=get_enum_type(scale_down_modes)) c.argument('allowed_host_ports', nargs='+', validator=validate_allowed_host_ports) c.argument('asg_ids', nargs='+', validator=validate_application_security_groups) + c.argument('enable_managed_dranet', action='store_true') c.argument('os_sku', arg_type=get_enum_type(node_os_skus_update), validator=validate_os_sku) c.argument("enable_fips_image", action="store_true") c.argument("disable_fips_image", action="store_true") diff --git a/src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py b/src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py index 70be49f2106..47f37a758cf 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py +++ b/src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py @@ -1786,6 +1786,10 @@ def get_allowed_host_ports(self) -> Union[List[PortRange], None]: )) return port_ranges + def get_enable_managed_dranet(self) -> bool: + """Obtain the value of enable_managed_dranet.""" + return self.raw_param.get("enable_managed_dranet", False) + def get_ip_tags(self) -> Union[List[IPTag], None]: ip_tags = self.raw_param.get("node_public_ip_tags") res = [] @@ -2329,10 +2333,15 @@ def set_up_agentpool_network_profile(self, agentpool: AgentPool) -> AgentPool: asg_ids = self.context.get_asg_ids() allowed_host_ports = self.context.get_allowed_host_ports() - if allowed_host_ports is not None: + enable_managed_dranet = self.context.get_enable_managed_dranet() + if allowed_host_ports is not None or enable_managed_dranet: agentpool.network_profile = self.models.AgentPoolNetworkProfile() agentpool.network_profile.allowed_host_ports = allowed_host_ports agentpool.network_profile.application_security_groups = asg_ids + if enable_managed_dranet: + agentpool.network_profile.dranet = self.models.DRANETProfile( + mode="Managed" + ) ip_tags = self.context.get_ip_tags() if ip_tags: @@ -2866,12 +2875,21 @@ def update_network_profile(self, agentpool: AgentPool) -> AgentPool: asg_ids = self.context.get_asg_ids() allowed_host_ports = self.context.get_allowed_host_ports() - if (asg_ids or allowed_host_ports) and not agentpool.network_profile: + enable_managed_dranet = self.context.get_enable_managed_dranet() + if ( + asg_ids is not None or + allowed_host_ports is not None or + enable_managed_dranet + ) and not agentpool.network_profile: agentpool.network_profile = self.models.AgentPoolNetworkProfile() if asg_ids is not None: agentpool.network_profile.application_security_groups = asg_ids if allowed_host_ports is not None: agentpool.network_profile.allowed_host_ports = allowed_host_ports + if enable_managed_dranet: + agentpool.network_profile.dranet = self.models.DRANETProfile( + mode="Managed" + ) return agentpool def update_os_sku(self, agentpool: AgentPool) -> AgentPool: diff --git a/src/azure-cli/azure/cli/command_modules/acs/custom.py b/src/azure-cli/azure/cli/command_modules/acs/custom.py index a0008c954d2..02ec9df017d 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/custom.py +++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py @@ -3039,6 +3039,7 @@ def aks_agentpool_add( gpu_instance_profile=None, allowed_host_ports=None, asg_ids=None, + enable_managed_dranet=False, node_public_ip_tags=None, disable_windows_outbound_nat=False, workload_runtime=None, @@ -3104,6 +3105,7 @@ def aks_agentpool_update( aks_custom_headers=None, allowed_host_ports=None, asg_ids=None, + enable_managed_dranet=False, os_sku=None, enable_fips_image=False, disable_fips_image=False, diff --git a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_agentpool_decorator.py b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_agentpool_decorator.py index f4b98145f5f..6f1b3f76212 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_agentpool_decorator.py +++ b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_agentpool_decorator.py @@ -1480,6 +1480,34 @@ def common_get_enable_artifact_streaming(self): ctx_2.attach_agentpool(agentpool_2) self.assertEqual(ctx_2.get_enable_artifact_streaming(), None) + def common_get_enable_managed_dranet(self): + ctx_0 = AKSAgentPoolContext( + self.cmd, + AKSAgentPoolParamDict({}), + self.models, + DecoratorMode.CREATE, + self.agentpool_decorator_mode, + ) + self.assertEqual(ctx_0.get_enable_managed_dranet(), False) + + ctx_1 = AKSAgentPoolContext( + self.cmd, + AKSAgentPoolParamDict({"enable_managed_dranet": False}), + self.models, + DecoratorMode.CREATE, + self.agentpool_decorator_mode, + ) + self.assertEqual(ctx_1.get_enable_managed_dranet(), False) + + ctx_2 = AKSAgentPoolContext( + self.cmd, + AKSAgentPoolParamDict({"enable_managed_dranet": True}), + self.models, + DecoratorMode.UPDATE, + self.agentpool_decorator_mode, + ) + self.assertEqual(ctx_2.get_enable_managed_dranet(), True) + def common_get_disable_artifact_streaming(self): # default ctx_1 = AKSAgentPoolContext( @@ -2086,6 +2114,9 @@ def test_get_disable_fips_image(self): def test_get_enable_artifact_streaming(self): self.common_get_enable_artifact_streaming() + def test_get_enable_managed_dranet(self): + self.common_get_enable_managed_dranet() + def test_get_disable_artifact_streaming(self): self.common_get_disable_artifact_streaming() @@ -2291,6 +2322,9 @@ def test_get_enable_fips_image(self): def test_get_enable_artifact_streaming(self): self.common_get_enable_artifact_streaming() + def test_get_enable_managed_dranet(self): + self.common_get_enable_managed_dranet() + def test_get_disable_artifact_streaming(self): self.common_get_disable_artifact_streaming() @@ -2884,6 +2918,31 @@ def common_set_up_agentpool_gateway_profile(self): ) self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1) + def common_set_up_managed_dranet(self): + dec_1 = AKSAgentPoolAddDecorator( + self.cmd, + self.client, + {"enable_managed_dranet": False}, + self.resource_type, + self.agentpool_decorator_mode, + ) + agentpool_1 = self.create_initialized_agentpool_instance(restore_defaults=False) + dec_1.context.attach_agentpool(agentpool_1) + dec_agentpool_1 = dec_1.set_up_agentpool_network_profile(agentpool_1) + self.assertIsNone(dec_agentpool_1.network_profile) + + dec_2 = AKSAgentPoolAddDecorator( + self.cmd, + self.client, + {"enable_managed_dranet": True}, + self.resource_type, + self.agentpool_decorator_mode, + ) + agentpool_2 = self.create_initialized_agentpool_instance(restore_defaults=False) + dec_2.context.attach_agentpool(agentpool_2) + dec_agentpool_2 = dec_2.set_up_agentpool_network_profile(agentpool_2) + self.assertEqual(dec_agentpool_2.network_profile.dranet.mode, "Managed") + def common_set_up_virtual_machines_profile(self): dec_1 = AKSAgentPoolAddDecorator( self.cmd, @@ -3155,6 +3214,9 @@ def test_set_up_gpu_propertes(self): def test_set_up_agentpool_gateway_profile(self): self.common_set_up_agentpool_gateway_profile() + def test_set_up_managed_dranet(self): + self.common_set_up_managed_dranet() + class AKSAgentPoolAddDecoratorManagedClusterModeTestCase(AKSAgentPoolAddDecoratorCommonTestCase): def setUp(self): self.cli_ctx = MockCLI() @@ -3283,6 +3345,9 @@ def test_set_up_gpu_propertes(self): def test_set_up_agentpool_gateway_profile(self): self.common_set_up_agentpool_gateway_profile() + def test_set_up_managed_dranet(self): + self.common_set_up_managed_dranet() + class AKSAgentPoolUpdateDecoratorCommonTestCase(unittest.TestCase): def _remove_defaults_in_agentpool(self, agentpool): self.defaults_in_agentpool = {} @@ -3653,6 +3718,52 @@ def common_update_artifact_streaming(self): with self.assertRaises(MutuallyExclusiveArgumentError): dec_3.update_artifact_streaming(agentpool_2) + def common_update_managed_dranet(self): + dec_1 = AKSAgentPoolUpdateDecorator( + self.cmd, + self.client, + {"enable_managed_dranet": False}, + self.resource_type, + self.agentpool_decorator_mode, + ) + agentpool_1 = self.create_initialized_agentpool_instance( + network_profile=self.models.AgentPoolNetworkProfile( + dranet=self.models.DRANETProfile(mode="Managed") + ) + ) + dec_1.context.attach_agentpool(agentpool_1) + dec_agentpool_1 = dec_1.update_network_profile(agentpool_1) + self.assertEqual(dec_agentpool_1.network_profile.dranet.mode, "Managed") + + dec_2 = AKSAgentPoolUpdateDecorator( + self.cmd, + self.client, + {"enable_managed_dranet": True}, + self.resource_type, + self.agentpool_decorator_mode, + ) + agentpool_2 = self.create_initialized_agentpool_instance() + dec_2.context.attach_agentpool(agentpool_2) + dec_agentpool_2 = dec_2.update_network_profile(agentpool_2) + self.assertEqual(dec_agentpool_2.network_profile.dranet.mode, "Managed") + + dec_3 = AKSAgentPoolUpdateDecorator( + self.cmd, + self.client, + {"enable_managed_dranet": False}, + self.resource_type, + self.agentpool_decorator_mode, + ) + agentpool_3 = self.create_initialized_agentpool_instance() + dec_3.context.attach_agentpool(agentpool_3) + with patch.object(dec_3.context, "get_asg_ids", return_value=[]), patch.object( + dec_3.context, "get_allowed_host_ports", return_value=[] + ): + dec_agentpool_3 = dec_3.update_network_profile(agentpool_3) + self.assertIsNotNone(dec_agentpool_3.network_profile) + self.assertEqual(dec_agentpool_3.network_profile.application_security_groups, []) + self.assertEqual(dec_agentpool_3.network_profile.allowed_host_ports, []) + def common_update_fips_image(self): dec_1 = AKSAgentPoolUpdateDecorator( self.cmd, @@ -3902,6 +4013,9 @@ def test_update_gpu_profile(self): def test_update_artifact_streaming(self): self.common_update_artifact_streaming() + def test_update_managed_dranet(self): + self.common_update_managed_dranet() + def test_update_agentpool_profile_default(self): import inspect @@ -4036,6 +4150,9 @@ def test_update_fips_image(self): def test_update_artifact_streaming(self): self.common_update_artifact_streaming() + def test_update_managed_dranet(self): + self.common_update_managed_dranet() + def test_update_agentpool_profile_default(self): import inspect