Skip to content

Commit 2586119

Browse files
author
Kanan Mehta
committed
feat: update validators and test
1 parent 9aedade commit 2586119

5 files changed

Lines changed: 152 additions & 3 deletions

File tree

‎src/aks-preview/azext_aks_preview/_validators.py‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -954,10 +954,20 @@ def validate_asm_egress_name(namespace):
954954

955955

956956
def validate_artifact_streaming(namespace):
957-
"""Validates that artifact streaming enablement can only be used on Linux."""
958-
if namespace.enable_artifact_streaming:
959-
if hasattr(namespace, 'os_type') and str(namespace.os_type).lower() == "windows":
957+
"""Validates artifact streaming flags for mutual exclusivity and OS support."""
958+
enable_artifact_streaming = getattr(namespace, "enable_artifact_streaming", False)
959+
disable_artifact_streaming = getattr(namespace, "disable_artifact_streaming", False)
960+
961+
if enable_artifact_streaming and disable_artifact_streaming:
962+
raise MutuallyExclusiveArgumentError(
963+
"Cannot specify both --enable-artifact-streaming and --disable-artifact-streaming at the same time."
964+
)
965+
966+
if hasattr(namespace, "os_type") and str(namespace.os_type).lower() == "windows":
967+
if enable_artifact_streaming:
960968
raise ArgumentUsageError('--enable-artifact-streaming can only be set for Linux nodepools')
969+
if disable_artifact_streaming:
970+
raise ArgumentUsageError('--disable-artifact-streaming can only be set for Linux nodepools')
961971

962972

963973
def validate_custom_endpoints(namespace):

‎src/aks-preview/azext_aks_preview/agentpool_decorator.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -584,8 +584,20 @@ def get_enable_artifact_streaming(self) -> bool:
584584
self.agentpool.artifact_streaming_profile.enabled is not None
585585
):
586586
enable_artifact_streaming = self.agentpool.artifact_streaming_profile.enabled
587+
588+
if enable_artifact_streaming and self.get_disable_artifact_streaming():
589+
raise MutuallyExclusiveArgumentError(
590+
'Cannot specify "--enable-artifact-streaming" and "--disable-artifact-streaming" at the same time'
591+
)
587592
return enable_artifact_streaming
588593

594+
def get_disable_artifact_streaming(self) -> bool:
595+
"""Obtain the value of disable_artifact_streaming.
596+
:return: bool
597+
"""
598+
599+
return self.raw_param.get("disable_artifact_streaming")
600+
589601
def get_pod_ip_allocation_mode(self: bool = False) -> Union[str, None]:
590602
"""Get the value of pod_ip_allocation_mode.
591603
:return: str or None
@@ -1273,6 +1285,13 @@ def set_up_artifact_streaming(self, agentpool: AgentPool) -> AgentPool:
12731285
self.models.AgentPoolArtifactStreamingProfile() # pylint: disable=no-member
12741286
)
12751287
agentpool.artifact_streaming_profile.enabled = True
1288+
1289+
if self.context.get_disable_artifact_streaming():
1290+
if agentpool.artifact_streaming_profile is None:
1291+
agentpool.artifact_streaming_profile = (
1292+
self.models.AgentPoolArtifactStreamingProfile() # pylint: disable=no-member
1293+
)
1294+
agentpool.artifact_streaming_profile.enabled = False
12761295
return agentpool
12771296

12781297
def set_up_ssh_access(self, agentpool: AgentPool) -> AgentPool:
@@ -1665,6 +1684,11 @@ def update_artifact_streaming(self, agentpool: AgentPool) -> AgentPool:
16651684
if agentpool.artifact_streaming_profile is None:
16661685
agentpool.artifact_streaming_profile = self.models.AgentPoolArtifactStreamingProfile() # pylint: disable=no-member
16671686
agentpool.artifact_streaming_profile.enabled = True
1687+
1688+
if self.context.get_disable_artifact_streaming():
1689+
if agentpool.artifact_streaming_profile is None:
1690+
agentpool.artifact_streaming_profile = self.models.AgentPoolArtifactStreamingProfile() # pylint: disable=no-member
1691+
agentpool.artifact_streaming_profile.enabled = False
16681692
return agentpool
16691693

16701694
def update_os_sku(self, agentpool: AgentPool) -> AgentPool:

‎src/aks-preview/azext_aks_preview/custom.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1893,6 +1893,7 @@ def aks_agentpool_add(
18931893
asg_ids=None,
18941894
node_public_ip_tags=None,
18951895
enable_artifact_streaming=False,
1896+
disable_artifact_streaming=False,
18961897
skip_gpu_driver_install=False,
18971898
gpu_driver=None,
18981899
driver_type=None,
@@ -1967,6 +1968,7 @@ def aks_agentpool_update(
19671968
allowed_host_ports=None,
19681969
asg_ids=None,
19691970
enable_artifact_streaming=False,
1971+
disable_artifact_streaming=False,
19701972
os_sku=None,
19711973
ssh_access=None,
19721974
yes=False,

‎src/aks-preview/azext_aks_preview/tests/latest/test_agentpool_decorator.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -256,6 +256,24 @@ def common_get_enable_artifact_streaming(self):
256256
ctx_2.attach_agentpool(agentpool_2)
257257
self.assertEqual(ctx_2.get_enable_artifact_streaming(), None)
258258

259+
def common_get_disable_artifact_streaming(self):
260+
# default
261+
ctx_1 = AKSPreviewAgentPoolContext(
262+
self.cmd,
263+
AKSAgentPoolParamDict({"disable_artifact_streaming": True}),
264+
self.models,
265+
DecoratorMode.UPDATE,
266+
self.agentpool_decorator_mode,
267+
)
268+
self.assertEqual(ctx_1.get_disable_artifact_streaming(), True)
269+
agentpool_1 = self.create_initialized_agentpool_instance(
270+
artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile(
271+
enabled=True
272+
)
273+
)
274+
ctx_1.attach_agentpool(agentpool_1)
275+
self.assertEqual(ctx_1.get_disable_artifact_streaming(), True)
276+
259277
def common_get_pod_ip_allocation_mode(self):
260278
# default
261279
ctx_1 = AKSPreviewAgentPoolContext(
@@ -1036,6 +1054,9 @@ def test_get_workload_runtime(self):
10361054
def test_get_enable_artifact_streaming(self):
10371055
self.common_get_enable_artifact_streaming()
10381056

1057+
def test_get_disable_artifact_streaming(self):
1058+
self.common_get_disable_artifact_streaming()
1059+
10391060
def test_get_pod_ip_allocation_mode(self):
10401061
self.common_get_pod_ip_allocation_mode()
10411062

@@ -1139,6 +1160,9 @@ def test_get_os_sku(self):
11391160
def test_get_enable_artifact_streaming(self):
11401161
self.common_get_enable_artifact_streaming()
11411162

1163+
def test_get_disable_artifact_streaming(self):
1164+
self.common_get_disable_artifact_streaming()
1165+
11421166
def test_get_enable_secure_boot(self):
11431167
self.common_get_enable_secure_boot()
11441168

@@ -2263,6 +2287,42 @@ def common_update_artifact_streaming(self):
22632287
)
22642288
self.assertEqual(dec_agentpool_2, grond_truth_agentpool_2)
22652289

2290+
dec_3 = AKSPreviewAgentPoolUpdateDecorator(
2291+
self.cmd,
2292+
self.client,
2293+
{"disable_artifact_streaming": True},
2294+
self.resource_type,
2295+
self.agentpool_decorator_mode,
2296+
)
2297+
# fail on passing the wrong agentpool object
2298+
with self.assertRaises(CLIInternalError):
2299+
dec_3.update_artifact_streaming(None)
2300+
agentpool_3 = self.create_initialized_agentpool_instance(
2301+
artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile(
2302+
enabled=True
2303+
)
2304+
)
2305+
dec_3.context.attach_agentpool(agentpool_3)
2306+
dec_agentpool_3 = dec_3.update_artifact_streaming(agentpool_3)
2307+
grond_truth_agentpool_3 = self.create_initialized_agentpool_instance(
2308+
artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile(
2309+
enabled=False
2310+
)
2311+
)
2312+
self.assertEqual(dec_agentpool_3, grond_truth_agentpool_3)
2313+
2314+
# Should error if both set
2315+
dec_4 = AKSPreviewAgentPoolUpdateDecorator(
2316+
self.cmd,
2317+
self.client,
2318+
{"enable_artifact_streaming": True, "disable_artifact_streaming": True},
2319+
self.resource_type,
2320+
self.agentpool_decorator_mode,
2321+
)
2322+
dec_4.context.attach_agentpool(agentpool_3)
2323+
with self.assertRaises(MutuallyExclusiveArgumentError):
2324+
dec_4.update_artifact_streaming(agentpool_3)
2325+
22662326
def common_update_secure_boot(self):
22672327
dec_1 = AKSPreviewAgentPoolUpdateDecorator(
22682328
self.cmd,

‎src/aks-preview/azext_aks_preview/tests/latest/test_validators.py‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,13 @@ def __init__(self, os_type, disable_windows_outbound_nat):
149149
self.disable_windows_outbound_nat = disable_windows_outbound_nat
150150

151151

152+
class ArtifactStreamingNamespace:
153+
def __init__(self, os_type, enable_artifact_streaming=False, disable_artifact_streaming=False):
154+
self.os_type = os_type
155+
self.enable_artifact_streaming = enable_artifact_streaming
156+
self.disable_artifact_streaming = disable_artifact_streaming
157+
158+
152159
class TestMaxSurge(unittest.TestCase):
153160
def test_valid_cases(self):
154161
valid = ["5", "33%", "1", "100%"]
@@ -385,6 +392,52 @@ def test_fail_if_os_type_invalid(self):
385392
)
386393

387394

395+
class TestArtifactStreaming(unittest.TestCase):
396+
def test_valid_linux_enable(self):
397+
validators.validate_artifact_streaming(
398+
ArtifactStreamingNamespace("Linux", enable_artifact_streaming=True)
399+
)
400+
401+
def test_valid_linux_disable(self):
402+
validators.validate_artifact_streaming(
403+
ArtifactStreamingNamespace("Linux", disable_artifact_streaming=True)
404+
)
405+
406+
def test_fail_if_enable_and_disable_are_set(self):
407+
with self.assertRaises(MutuallyExclusiveArgumentError) as cm:
408+
validators.validate_artifact_streaming(
409+
ArtifactStreamingNamespace(
410+
"Linux",
411+
enable_artifact_streaming=True,
412+
disable_artifact_streaming=True,
413+
)
414+
)
415+
self.assertEqual(
416+
str(cm.exception),
417+
"Cannot specify both --enable-artifact-streaming and --disable-artifact-streaming at the same time.",
418+
)
419+
420+
def test_fail_if_enable_for_windows(self):
421+
with self.assertRaises(ArgumentUsageError) as cm:
422+
validators.validate_artifact_streaming(
423+
ArtifactStreamingNamespace("Windows", enable_artifact_streaming=True)
424+
)
425+
self.assertEqual(
426+
str(cm.exception),
427+
"--enable-artifact-streaming can only be set for Linux nodepools",
428+
)
429+
430+
def test_fail_if_disable_for_windows(self):
431+
with self.assertRaises(ArgumentUsageError) as cm:
432+
validators.validate_artifact_streaming(
433+
ArtifactStreamingNamespace("Windows", disable_artifact_streaming=True)
434+
)
435+
self.assertEqual(
436+
str(cm.exception),
437+
"--disable-artifact-streaming can only be set for Linux nodepools",
438+
)
439+
440+
388441
class ValidateAddonsNamespace:
389442
def __init__(self, addons):
390443
self.addons = addons

0 commit comments

Comments
 (0)