From b34fa7f7ecb1389b23a139c947096303a88dc69a Mon Sep 17 00:00:00 2001 From: Fuming Zhang Date: Wed, 18 Oct 2023 14:37:51 +0800 Subject: [PATCH] update --- .../acs/agentpool_decorator.py | 2 +- .../tests/latest/test_agentpool_decorator.py | 20 +++++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) 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 dee4a99e124..70c935e1e81 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 @@ -943,7 +943,7 @@ def get_drain_timeout(self): if ( self.agentpool and self.agentpool.upgrade_settings and - self.agentpool.upgrade_settings.drain_timeout is not None + self.agentpool.upgrade_settings.drain_timeout_in_minutes is not None ): drain_timeout = self.agentpool.upgrade_settings.drain_timeout_in_minutes 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 4464dbfa874..1603d20b294 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 @@ -1308,6 +1308,20 @@ def common_get_gpu_instance_profile(self): ctx_1.attach_agentpool(agentpool_1) self.assertEqual(ctx_1.get_gpu_instance_profile(), "test_gpu_instance_profile") + def common_get_drain_timeout(self): + # default + ctx_1 = AKSAgentPoolContext( + self.cmd, + AKSAgentPoolParamDict({"drain_timeout": None}), + self.models, + DecoratorMode.CREATE, + self.agentpool_decorator_mode, + ) + self.assertEqual(ctx_1.get_drain_timeout(), None) + agentpool_1 = self.create_initialized_agentpool_instance(upgrade_settings=self.models.AgentPoolUpgradeSettings(drain_timeout_in_minutes=123)) + ctx_1.attach_agentpool(agentpool_1) + self.assertEqual(ctx_1.get_drain_timeout(), 123) + class AKSAgentPoolContextStandaloneModeTestCase(AKSAgentPoolContextCommonTestCase): def setUp(self): self.cli_ctx = MockCLI() @@ -1488,6 +1502,9 @@ def test_get_no_wait(self): def test_get_gpu_instance_profile(self): self.common_get_gpu_instance_profile() + def test_get_drain_timeout(self): + self.common_get_drain_timeout() + class AKSAgentPoolContextManagedClusterModeTestCase(AKSAgentPoolContextCommonTestCase): def setUp(self): self.cli_ctx = MockCLI() @@ -1642,6 +1659,9 @@ def test_get_no_wait(self): def test_get_gpu_instance_profile(self): self.common_get_gpu_instance_profile() + def test_get_drain_timeout(self): + self.common_get_drain_timeout() + class AKSAgentPoolAddDecoratorCommonTestCase(unittest.TestCase): def _remove_defaults_in_agentpool(self, agentpool): self.defaults_in_agentpool = {}