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 7a89268ed9f..6b338fd5fd3 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/_help.py +++ b/src/azure-cli/azure/cli/command_modules/acs/_help.py @@ -2010,6 +2010,9 @@ long-summary: | Azure provides a different workload-runtime to enable Kata supported workloads in your nodepools. The following values can be specified: - "KataVmIsolation" for Kata. + - name: --enable-artifact-streaming + type: bool + short-summary: Enable artifact streaming for VirtualMachineScaleSets managed by a node pool, to speed up the cold-start of containers on a node through on-demand image loading. This option is only valid for Linux nodepools. To use this feature, container images must also enable artifact streaming on ACR. If not specified, the default is false. examples: - name: Create a nodepool in an existing AKS cluster with ephemeral os enabled. @@ -2168,6 +2171,12 @@ - name: --gpu-driver type: string short-summary: Whether to install driver for GPU node pool. Possible values are "Install" or "None". + - name: --enable-artifact-streaming + type: bool + short-summary: Enable artifact streaming for VirtualMachineScaleSets managed by a Linux node pool, to speed up the cold-start of containers on a node through on-demand image loading. To use this feature, container images must also enable artifact streaming on ACR. If not specified, the default is false. + - name: --disable-artifact-streaming + type: bool + short-summary: Disable artifact streaming for VirtualMachineScaleSets managed by a Linux node pool. examples: - name: Reconcile the nodepool back to its current state. text: az aks nodepool update -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster 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 49e2d17f585..0120800c2db 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/_params.py +++ b/src/azure-cli/azure/cli/command_modules/acs/_params.py @@ -138,6 +138,7 @@ validate_custom_ca_trust_certificates, validate_bootstrap_container_registry_resource_id, validate_gateway_prefix_size, + validate_artifact_streaming, ) from azure.cli.core.commands.parameters import ( edge_zone_type, file_type, get_enum_type, @@ -1095,6 +1096,7 @@ def load_arguments(self, _): c.argument("gateway_prefix_size", type=int, validator=validate_gateway_prefix_size) c.argument('localdns_config', help='Path to a JSON file to configure the local DNS profile for a new nodepool.') c.argument('workload_runtime', arg_type=get_enum_type(workload_runtime_types), help="The workload runtime to use on the nodepool.") + c.argument('enable_artifact_streaming', action='store_true', validator=validate_artifact_streaming) with self.argument_context('aks nodepool update', resource_type=ResourceType.MGMT_CONTAINERSERVICE, operation_group='agent_pools') as c: c.argument('enable_cluster_autoscaler', options_list=[ @@ -1128,6 +1130,8 @@ def load_arguments(self, _): c.argument("if_none_match") c.argument('localdns_config', help='Path to a JSON file to configure the local DNS profile for a new nodepool.') c.argument('gpu_driver', arg_type=get_enum_type(gpu_driver_install_modes)) + c.argument('enable_artifact_streaming', action='store_true', validator=validate_artifact_streaming) + c.argument('disable_artifact_streaming', action='store_true', validator=validate_artifact_streaming) with self.argument_context('aks nodepool upgrade') as c: c.argument('max_surge', validator=validate_max_surge) diff --git a/src/azure-cli/azure/cli/command_modules/acs/_validators.py b/src/azure-cli/azure/cli/command_modules/acs/_validators.py index 45ab88989d9..b1dbb1af1ac 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/_validators.py +++ b/src/azure-cli/azure/cli/command_modules/acs/_validators.py @@ -971,3 +971,20 @@ def validate_gateway_prefix_size(namespace): raise ArgumentUsageError("--gateway-prefix-size can only be set for Gateway-mode nodepools") if namespace.gateway_prefix_size < 28 or namespace.gateway_prefix_size > 31: raise InvalidArgumentValueError("--gateway-prefix-size must be in the range [28, 31]") + + +def validate_artifact_streaming(namespace): + """Validates artifact streaming flags for mutual exclusivity and OS support.""" + enable_artifact_streaming = getattr(namespace, "enable_artifact_streaming", False) + disable_artifact_streaming = getattr(namespace, "disable_artifact_streaming", False) + + if enable_artifact_streaming and disable_artifact_streaming: + raise MutuallyExclusiveArgumentError( + "Cannot specify both --enable-artifact-streaming and --disable-artifact-streaming at the same time." + ) + + if hasattr(namespace, "os_type") and str(namespace.os_type).lower() == "windows": + if enable_artifact_streaming: + raise ArgumentUsageError('--enable-artifact-streaming can only be set for Linux nodepools') + if disable_artifact_streaming: + raise ArgumentUsageError('--disable-artifact-streaming can only be set for Linux nodepools') 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 f9b46379bad..c7438d94903 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 @@ -1354,6 +1354,35 @@ 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_enable_artifact_streaming(self) -> bool: + """Obtain the value of enable_artifact_streaming. + + :return: bool + """ + # read the original value passed by the command + enable_artifact_streaming = self.raw_param.get("enable_artifact_streaming") + # In create mode, try to read the property value corresponding to the parameter from the `agentpool` object + if self.decorator_mode == DecoratorMode.CREATE: + if ( + self.agentpool and + self.agentpool.artifact_streaming_profile is not None and + self.agentpool.artifact_streaming_profile.enabled is not None + ): + enable_artifact_streaming = self.agentpool.artifact_streaming_profile.enabled + + if enable_artifact_streaming and self.get_disable_artifact_streaming(): + raise MutuallyExclusiveArgumentError( + 'Cannot specify both --enable-artifact-streaming and --disable-artifact-streaming.' + ) + return enable_artifact_streaming + + def get_disable_artifact_streaming(self) -> bool: + """Obtain the value of disable_artifact_streaming. + :return: bool + """ + + return self.raw_param.get("disable_artifact_streaming") + def get_zones(self) -> Union[List[str], None]: """Obtain the value of zones. @@ -2188,6 +2217,16 @@ def set_up_gpu_profile(self, agentpool: AgentPool) -> AgentPool: return agentpool + def set_up_artifact_streaming(self, agentpool: AgentPool) -> AgentPool: + """Set up artifact streaming property for the AgentPool object.""" + self._ensure_agentpool(agentpool) + + if self.context.get_enable_artifact_streaming(): + if agentpool.artifact_streaming_profile is None: + agentpool.artifact_streaming_profile = self.models.AgentPoolArtifactStreamingProfile() # pylint: disable=no-member + agentpool.artifact_streaming_profile.enabled = True + return agentpool + def set_up_agentpool_gateway_profile(self, agentpool: AgentPool) -> AgentPool: """Set up agentpool gateway profile for the AgentPool object. @@ -2297,6 +2336,8 @@ def construct_agentpool_profile_default(self, bypass_restore_defaults: bool = Fa agentpool = self.set_up_gpu_properties(agentpool) # set up agentpool network profile agentpool = self.set_up_agentpool_network_profile(agentpool) + # set up artifact streaming + agentpool = self.set_up_artifact_streaming(agentpool) # set up agentpool pod ip allocation mode agentpool = self.set_up_pod_ip_allocation_mode(agentpool) # set up agentpool windows profile @@ -2636,6 +2677,23 @@ def update_gpu_profile(self, agentpool: AgentPool) -> AgentPool: return agentpool + def update_artifact_streaming(self, agentpool: AgentPool) -> AgentPool: + """Update artifact streaming property for the AgentPool object. + :return: the AgentPool object + """ + self._ensure_agentpool(agentpool) + + if self.context.get_enable_artifact_streaming(): + if agentpool.artifact_streaming_profile is None: + agentpool.artifact_streaming_profile = self.models.AgentPoolArtifactStreamingProfile() # pylint: disable=no-member + agentpool.artifact_streaming_profile.enabled = True + + if self.context.get_disable_artifact_streaming(): + if agentpool.artifact_streaming_profile is None: + agentpool.artifact_streaming_profile = self.models.AgentPoolArtifactStreamingProfile() # pylint: disable=no-member + agentpool.artifact_streaming_profile.enabled = False + return agentpool + def update_agentpool_profile_default(self, agentpools: List[AgentPool] = None) -> AgentPool: """The overall controller used to update AgentPool profile by default. @@ -2668,6 +2726,8 @@ def update_agentpool_profile_default(self, agentpools: List[AgentPool] = None) - agentpool = self.update_localdns_profile(agentpool) # update gpu profile agentpool = self.update_gpu_profile(agentpool) + # update artifact streaming + agentpool = self.update_artifact_streaming(agentpool) return agentpool def update_agentpool(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 b0d7ece1360..72b82319f87 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/custom.py +++ b/src/azure-cli/azure/cli/command_modules/acs/custom.py @@ -2988,6 +2988,7 @@ def aks_agentpool_add( gateway_prefix_size=None, # local DNS localdns_config=None, + enable_artifact_streaming=False, ): # DO NOT MOVE: get all the original parameters and save them as a dictionary raw_parameters = locals() @@ -3051,6 +3052,8 @@ def aks_agentpool_update( # local DNS localdns_config=None, gpu_driver=None, + enable_artifact_streaming=False, + disable_artifact_streaming=False, ): # DO NOT MOVE: get all the original parameters and save them as a dictionary raw_parameters = locals() diff --git a/src/azure-cli/azure/cli/command_modules/acs/linter_exclusions.yml b/src/azure-cli/azure/cli/command_modules/acs/linter_exclusions.yml index 39692b2f608..cec412e4f51 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/linter_exclusions.yml +++ b/src/azure-cli/azure/cli/command_modules/acs/linter_exclusions.yml @@ -244,11 +244,20 @@ aks nodepool add: undrainable_node_behavior: rule_exclusions: - option_length_too_long + enable_artifact_streaming: + rule_exclusions: + - option_length_too_long aks nodepool update: parameters: undrainable_node_behavior: rule_exclusions: - option_length_too_long + enable_artifact_streaming: + rule_exclusions: + - option_length_too_long + disable_artifact_streaming: + rule_exclusions: + - option_length_too_long aks nodepool upgrade: parameters: undrainable_node_behavior: 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 0a00f71ffae..fa64e83a37d 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 @@ -1352,6 +1352,59 @@ def common_get_disable_fips_image(self): ctx_1.attach_agentpool(agentpool_1) self.assertEqual(ctx_1.get_disable_fips_image(), True) + def common_get_enable_artifact_streaming(self): + # default + ctx_1 = AKSAgentPoolContext( + self.cmd, + AKSAgentPoolParamDict({"enable_artifact_streaming": None}), + self.models, + DecoratorMode.CREATE, + self.agentpool_decorator_mode, + ) + self.assertEqual(ctx_1.get_enable_artifact_streaming(), None) + agentpool_1 = self.create_initialized_agentpool_instance( + artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile( + enabled=True + ) + ) + ctx_1.attach_agentpool(agentpool_1) + self.assertEqual(ctx_1.get_enable_artifact_streaming(), True) + + # update mode: do not read from agentpool property + ctx_2 = AKSAgentPoolContext( + self.cmd, + AKSAgentPoolParamDict({"enable_artifact_streaming": None}), + self.models, + DecoratorMode.UPDATE, + self.agentpool_decorator_mode, + ) + self.assertEqual(ctx_2.get_enable_artifact_streaming(), None) + agentpool_2 = self.create_initialized_agentpool_instance( + artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile( + enabled=True + ) + ) + ctx_2.attach_agentpool(agentpool_2) + self.assertEqual(ctx_2.get_enable_artifact_streaming(), None) + + def common_get_disable_artifact_streaming(self): + # default + ctx_1 = AKSAgentPoolContext( + self.cmd, + AKSAgentPoolParamDict({"disable_artifact_streaming": True}), + self.models, + DecoratorMode.UPDATE, + self.agentpool_decorator_mode, + ) + self.assertEqual(ctx_1.get_disable_artifact_streaming(), True) + agentpool_1 = self.create_initialized_agentpool_instance( + artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile( + enabled=True + ) + ) + ctx_1.attach_agentpool(agentpool_1) + self.assertEqual(ctx_1.get_disable_artifact_streaming(), True) + def common_get_zones(self): # default ctx_1 = AKSAgentPoolContext( @@ -1931,6 +1984,12 @@ def test_get_enable_fips_image(self): def test_get_disable_fips_image(self): self.common_get_disable_fips_image() + def test_get_enable_artifact_streaming(self): + self.common_get_enable_artifact_streaming() + + def test_get_disable_artifact_streaming(self): + self.common_get_disable_artifact_streaming() + def test_get_zones(self): self.common_get_zones() @@ -2124,6 +2183,12 @@ def test_get_enable_ultra_ssd(self): def test_get_enable_fips_image(self): self.common_get_enable_fips_image() + def test_get_enable_artifact_streaming(self): + self.common_get_enable_artifact_streaming() + + def test_get_disable_artifact_streaming(self): + self.common_get_disable_artifact_streaming() + def test_get_zones(self): self.common_get_zones() @@ -2670,6 +2735,28 @@ def common_set_up_gpu_profile(self): ) self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1) + def common_set_up_artifact_streaming(self): + dec_1 = AKSAgentPoolAddDecorator( + self.cmd, + self.client, + {"enable_artifact_streaming": True}, + self.resource_type, + self.agentpool_decorator_mode, + ) + # fail on passing the wrong agentpool object + with self.assertRaises(CLIInternalError): + dec_1.set_up_artifact_streaming(None) + agentpool_1 = self.create_initialized_agentpool_instance(restore_defaults=False) + dec_1.context.attach_agentpool(agentpool_1) + dec_agentpool_1 = dec_1.set_up_artifact_streaming(agentpool_1) + dec_agentpool_1 = self._restore_defaults_in_agentpool(dec_agentpool_1) + ground_truth_agentpool_1 = self.create_initialized_agentpool_instance( + artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile( + enabled=True + ) + ) + self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1) + def common_set_up_agentpool_gateway_profile(self): dec_1 = AKSAgentPoolAddDecorator( self.cmd, @@ -2805,6 +2892,9 @@ def test_set_up_agentpool_security_profile(self): def test_set_up_gpu_profile(self): self.common_set_up_gpu_profile() + def test_set_up_artifact_streaming(self): + self.common_set_up_artifact_streaming() + def test_set_up_virtual_machines_profile(self): self.common_set_up_virtual_machines_profile() @@ -3275,6 +3365,64 @@ def common_update_gpu_profile(self): ) self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1) + def common_update_artifact_streaming(self): + # enable + dec_1 = AKSAgentPoolUpdateDecorator( + self.cmd, + self.client, + {"enable_artifact_streaming": True, "disable_artifact_streaming": False}, + self.resource_type, + self.agentpool_decorator_mode, + ) + # fail on passing the wrong agentpool object + with self.assertRaises(CLIInternalError): + dec_1.update_artifact_streaming(None) + agentpool_1 = self.create_initialized_agentpool_instance() + dec_1.context.attach_agentpool(agentpool_1) + dec_agentpool_1 = dec_1.update_artifact_streaming(agentpool_1) + ground_truth_agentpool_1 = self.create_initialized_agentpool_instance( + artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile( + enabled=True + ) + ) + self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1) + + # disable + dec_2 = AKSAgentPoolUpdateDecorator( + self.cmd, + self.client, + {"enable_artifact_streaming": False, "disable_artifact_streaming": True}, + self.resource_type, + self.agentpool_decorator_mode, + ) + with self.assertRaises(CLIInternalError): + dec_2.update_artifact_streaming(None) + agentpool_2 = self.create_initialized_agentpool_instance( + artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile( + enabled=True + ) + ) + dec_2.context.attach_agentpool(agentpool_2) + dec_agentpool_2 = dec_2.update_artifact_streaming(agentpool_2) + ground_truth_agentpool_2 = self.create_initialized_agentpool_instance( + artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile( + enabled=False + ) + ) + self.assertEqual(dec_agentpool_2, ground_truth_agentpool_2) + + # Should error if both set + dec_3 = AKSAgentPoolUpdateDecorator( + self.cmd, + self.client, + {"enable_artifact_streaming": True, "disable_artifact_streaming": True}, + self.resource_type, + self.agentpool_decorator_mode, + ) + dec_3.context.attach_agentpool(agentpool_2) + with self.assertRaises(MutuallyExclusiveArgumentError): + dec_3.update_artifact_streaming(agentpool_2) + def common_update_fips_image(self): dec_1 = AKSAgentPoolUpdateDecorator( self.cmd, @@ -3518,6 +3666,9 @@ def test_update_fips_image(self): def test_update_gpu_profile(self): self.common_update_gpu_profile() + def test_update_artifact_streaming(self): + self.common_update_artifact_streaming() + def test_update_agentpool_profile_default(self): import inspect @@ -3646,6 +3797,9 @@ def test_update_upgrade_settings(self): def test_update_fips_image(self): self.common_update_fips_image() + def test_update_artifact_streaming(self): + self.common_update_artifact_streaming() + def test_update_agentpool_profile_default(self): import inspect diff --git a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_aks_commands.py b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_aks_commands.py index 93ee3ef724c..84ad09fc1f4 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_aks_commands.py +++ b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_aks_commands.py @@ -4088,6 +4088,134 @@ def test_aks_nodepool_add_with_disable_windows_outbound_nat( checks=[self.is_empty()], ) + # the availability of features is controlled by a toggle and cannot be fully tested yet, + # however, existing test results show that the client side works as expected, so exclude it at this moment + @live_only() + @AllowLargeResponse() + @AKSCustomResourceGroupPreparer( + random_name_length=17, name_prefix="clitest", location="eastus" + ) + def test_aks_nodepool_add_with_artifact_streaming( + self, resource_group, resource_group_location + ): + # reset the count so in replay mode the random names will start with 0 + self.test_resources_count = 0 + # kwargs for string formatting + aks_name = self.create_random_name("cliakstest", 16) + self.kwargs.update( + { + "resource_group": resource_group, + "name": aks_name, + "dns_name_prefix": self.create_random_name("cliaksdns", 16), + "location": resource_group_location, + "resource_type": "Microsoft.ContainerService/ManagedClusters", + "nodepool2_name": "np2", + "ssh_key_value": self.generate_ssh_keys(), + } + ) + + # create + create_cmd = ( + "aks create --resource-group={resource_group} --name={name} " + "--ssh-key-value={ssh_key_value} " + "--aks-custom-headers=AKSHTTPCustomFeatures=Microsoft.ContainerService/ArtifactStreamingPreview " + ) + self.cmd( + create_cmd, + checks=[ + self.check("provisioningState", "Succeeded"), + ], + ) + + # nodepool add + self.cmd( + "aks nodepool add --resource-group={resource_group} --cluster-name={name} --name={nodepool2_name} " + "--enable-artifact-streaming --aks-custom-headers=AKSHTTPCustomFeatures=Microsoft.ContainerService/ArtifactStreamingPreview", + checks=[ + self.check("provisioningState", "Succeeded"), + self.check("artifactStreamingProfile.enabled", True), + ], + ) + + # delete + self.cmd( + "aks delete -g {resource_group} -n {name} --yes --no-wait", + checks=[self.is_empty()], + ) + + # the availability of features is controlled by a toggle and cannot be fully tested yet, + # however, existing test results show that the client side works as expected, so exclude it at this moment + @live_only() + @AllowLargeResponse() + @AKSCustomResourceGroupPreparer( + random_name_length=17, name_prefix="clitest", location="eastus" + ) + def test_aks_nodepool_update_with_artifact_streaming( + self, resource_group, resource_group_location + ): + aks_name = self.create_random_name("cliakstest", 16) + nodepool_name = self.create_random_name("n", 6) + + self.kwargs.update( + { + "resource_group": resource_group, + "name": aks_name, + "location": resource_group_location, + "ssh_key_value": self.generate_ssh_keys(), + "node_pool_name": nodepool_name, + "node_vm_size": "standard_d2s_v3", + } + ) + + self.cmd( + "aks create " + "--resource-group={resource_group} " + "--name={name} " + "--location={location} " + "--ssh-key-value={ssh_key_value} " + "--nodepool-name={node_pool_name} " + "--node-count=1 " + "--node-vm-size={node_vm_size} " + "--aks-custom-headers AKSHTTPCustomFeatures=Microsoft.ContainerService/ArtifactStreamingPreview", + checks=[ + self.check("provisioningState", "Succeeded"), + ], + ) + + # enable artifact streaming + self.cmd( + "aks nodepool update " + "--resource-group={resource_group} " + "--cluster-name={name} " + "--name={node_pool_name} " + "--enable-artifact-streaming " + "--aks-custom-headers AKSHTTPCustomFeatures=Microsoft.ContainerService/ArtifactStreamingPreview", + checks=[ + self.check("provisioningState", "Succeeded"), + self.check("artifactStreamingProfile.enabled", True), + ], + ) + + # disable artifact streaming + self.cmd( + "aks nodepool update " + "--resource-group={resource_group} " + "--cluster-name={name} " + "--name={node_pool_name} " + "--disable-artifact-streaming " + "--aks-custom-headers AKSHTTPCustomFeatures=Microsoft.ContainerService/ArtifactStreamingPreview", + checks=[ + self.check("provisioningState", "Succeeded"), + self.check("artifactStreamingProfile.enabled", False), + ], + ) + + # delete + self.cmd( + "aks delete --resource-group={resource_group} --name={name} --yes --no-wait", + checks=[self.is_empty()], + ) + @AllowLargeResponse() @AKSCustomResourceGroupPreparer(random_name_length=17, name_prefix='clitest', location='westcentralus') def test_aks_control_plane_user_assigned_identity(self, resource_group, resource_group_location): diff --git a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_validators.py b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_validators.py index 52a54e5c802..6dc41fe4f53 100644 --- a/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_validators.py +++ b/src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_validators.py @@ -1933,5 +1933,58 @@ def test_fail_if_os_type_invalid(self): self.assertTrue('--custom-ca-trust-certificates can only be set for linux nodepools' in str(cm.exception), msg=str(cm.exception)) +class ArtifactStreamingNamespace: + def __init__(self, os_type, enable_artifact_streaming=False, disable_artifact_streaming=False): + self.os_type = os_type + self.enable_artifact_streaming = enable_artifact_streaming + self.disable_artifact_streaming = disable_artifact_streaming + + +class TestArtifactStreaming(unittest.TestCase): + def test_valid_linux_enable(self): + validators.validate_artifact_streaming( + ArtifactStreamingNamespace("Linux", enable_artifact_streaming=True) + ) + + def test_valid_linux_disable(self): + validators.validate_artifact_streaming( + ArtifactStreamingNamespace("Linux", disable_artifact_streaming=True) + ) + + def test_fail_if_enable_and_disable_are_set(self): + with self.assertRaises(MutuallyExclusiveArgumentError) as cm: + validators.validate_artifact_streaming( + ArtifactStreamingNamespace( + "Linux", + enable_artifact_streaming=True, + disable_artifact_streaming=True, + ) + ) + self.assertEqual( + str(cm.exception), + "Cannot specify both --enable-artifact-streaming and --disable-artifact-streaming at the same time.", + ) + + def test_fail_if_enable_for_windows(self): + with self.assertRaises(ArgumentUsageError) as cm: + validators.validate_artifact_streaming( + ArtifactStreamingNamespace("Windows", enable_artifact_streaming=True) + ) + self.assertEqual( + str(cm.exception), + "--enable-artifact-streaming can only be set for Linux nodepools", + ) + + def test_fail_if_disable_for_windows(self): + with self.assertRaises(ArgumentUsageError) as cm: + validators.validate_artifact_streaming( + ArtifactStreamingNamespace("Windows", disable_artifact_streaming=True) + ) + self.assertEqual( + str(cm.exception), + "--disable-artifact-streaming can only be set for Linux nodepools", + ) + + if __name__ == "__main__": unittest.main()