fix: address failing tests

Author: CustardTart32

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -498,10 +498,9 @@ def get_upgrade_strategy(self) -> Union[str, None]:
         if self.decorator_mode == DecoratorMode.CREATE:
             if (
                 self.agentpool and
-                self.agentpool.upgrade_settings and
-                self.agentpool.upgrade_settings.upgrade_strategy is not None
+                self.agentpool.upgrade_strategy is not None
             ):
-                upgrade_strategy = self.agentpool.upgrade_settings.upgrade_strategy
+                upgrade_strategy = self.agentpool.upgrade_strategy
 
         # this parameter does not need dynamic completion
         # this parameter does not need validation
@@ -518,9 +517,7 @@ def get_drain_batch_size(self) -> Union[str, None]:
         if self.decorator_mode == DecoratorMode.CREATE:
             if (
                 self.agentpool and
-                hasattr(self.agentpool, 'upgrade_settings_blue_green') and
-                self.agentpool.upgrade_settings_blue_green and
-                hasattr(self.agentpool.upgrade_settings_blue_green, 'drain_batch_size') and
+                self.agentpool.upgrade_settings_blue_green is not None and
                 self.agentpool.upgrade_settings_blue_green.drain_batch_size is not None
             ):
                 drain_batch_size = self.agentpool.upgrade_settings_blue_green.drain_batch_size
@@ -540,9 +537,7 @@ def get_drain_timeout_bg(self) -> Union[int, None]:
         if self.decorator_mode == DecoratorMode.CREATE:
             if (
                 self.agentpool and
-                hasattr(self.agentpool, 'upgrade_settings_blue_green') and
-                self.agentpool.upgrade_settings_blue_green and
-                hasattr(self.agentpool.upgrade_settings_blue_green, 'drain_timeout_in_minutes') and
+                self.agentpool.upgrade_settings_blue_green is not None and
                 self.agentpool.upgrade_settings_blue_green.drain_timeout_in_minutes is not None
             ):
                 drain_timeout_bg = self.agentpool.upgrade_settings_blue_green.drain_timeout_in_minutes
@@ -562,9 +557,7 @@ def get_batch_soak_duration(self) -> Union[int, None]:
         if self.decorator_mode == DecoratorMode.CREATE:
             if (
                 self.agentpool and
-                hasattr(self.agentpool, 'upgrade_settings_blue_green') and
-                self.agentpool.upgrade_settings_blue_green and
-                hasattr(self.agentpool.upgrade_settings_blue_green, 'batch_soak_duration_in_minutes') and
+                self.agentpool.upgrade_settings_blue_green is not None and
                 self.agentpool.upgrade_settings_blue_green.batch_soak_duration_in_minutes is not None
             ):
                 batch_soak_duration = self.agentpool.upgrade_settings_blue_green.batch_soak_duration_in_minutes
@@ -584,9 +577,7 @@ def get_final_soak_duration(self) -> Union[int, None]:
         if self.decorator_mode == DecoratorMode.CREATE:
             if (
                 self.agentpool and
-                hasattr(self.agentpool, 'upgrade_settings_blue_green') and
-                self.agentpool.upgrade_settings_blue_green and
-                hasattr(self.agentpool.upgrade_settings_blue_green, 'final_soak_duration_in_minutes') and
+                self.agentpool.upgrade_settings_blue_green is not None and
                 self.agentpool.upgrade_settings_blue_green.final_soak_duration_in_minutes is not None
             ):
                 final_soak_duration = self.agentpool.upgrade_settings_blue_green.final_soak_duration_in_minutes
@@ -1361,37 +1352,6 @@ def set_up_upgrade_strategy(self, agentpool: AgentPool) -> AgentPool:
 
         return agentpool
 
-    def set_up_blue_green_upgrade_settings(self, agentpool: AgentPool) -> AgentPool:
-        """Set up blue-green upgrade settings for the AgentPool object.
-
-        :return: the AgentPool object
-        """
-        self._ensure_agentpool(agentpool)
-        
-        blue_green_upgrade_settings = self.models.AgentPoolBlueGreenUpgradeSettings()
-            
-        # Set each field if provided in context
-        drain_batch_size = self.context.get_drain_batch_size()
-        if drain_batch_size is not None:
-            blue_green_upgrade_settings.drain_batch_size = drain_batch_size
-            
-        drain_timeout_bg = self.context.get_drain_timeout_bg()
-        if drain_timeout_bg is not None:
-            blue_green_upgrade_settings.drain_timeout_in_minutes = drain_timeout_bg
-        
-        batch_soak_duration = self.context.get_batch_soak_duration()
-        if batch_soak_duration is not None:
-            blue_green_upgrade_settings.batch_soak_duration_in_minutes = batch_soak_duration
-            
-        final_soak_duration = self.context.get_final_soak_duration()
-        if final_soak_duration is not None:
-            blue_green_upgrade_settings.final_soak_duration_in_minutes = final_soak_duration
-            
-        # Set the blue-green upgrade settings as a separate property on agentpool
-        agentpool.upgrade_settings_blue_green = blue_green_upgrade_settings
-
-        return agentpool
-
     def set_up_upgrade_settings(self, agentpool: AgentPool) -> AgentPool:
         """Set up upgrade settings for the AgentPool object.
 
@@ -1427,6 +1387,36 @@ def set_up_upgrade_settings(self, agentpool: AgentPool) -> AgentPool:
         agentpool.upgrade_settings = upgrade_settings
         return agentpool
 
+    def set_up_blue_green_upgrade_settings(self, agentpool: AgentPool) -> AgentPool:
+        """Set up blue-green upgrade settings for the AgentPool object.
+
+        :return: the AgentPool object
+        """
+        self._ensure_agentpool(agentpool)
+        
+        blue_green_upgrade_settings = self.models.AgentPoolBlueGreenUpgradeSettings()
+            
+        # Set each field if provided in context
+        drain_batch_size = self.context.get_drain_batch_size()
+        if drain_batch_size is not None:
+            blue_green_upgrade_settings.drain_batch_size = drain_batch_size
+            
+        drain_timeout_bg = self.context.get_drain_timeout_bg()
+        if drain_timeout_bg is not None:
+            blue_green_upgrade_settings.drain_timeout_in_minutes = drain_timeout_bg
+        
+        batch_soak_duration = self.context.get_batch_soak_duration()
+        if batch_soak_duration is not None:
+            blue_green_upgrade_settings.batch_soak_duration_in_minutes = batch_soak_duration
+            
+        final_soak_duration = self.context.get_final_soak_duration()
+        if final_soak_duration is not None:
+            blue_green_upgrade_settings.final_soak_duration_in_minutes = final_soak_duration
+            
+        # Set the blue-green upgrade settings as a separate property on agentpool
+        agentpool.upgrade_settings_blue_green = blue_green_upgrade_settings
+
+        return agentpool
 
 class AKSPreviewAgentPoolUpdateDecorator(AKSAgentPoolUpdateDecorator):
     def __init__(
@@ -1746,27 +1736,28 @@ def update_blue_green_upgrade_settings(self, agentpool: AgentPool) -> AgentPool:
         drain_timeout_bg = self.context.get_drain_timeout_bg()
         batch_soak_duration = self.context.get_batch_soak_duration()
         final_soak_duration = self.context.get_final_soak_duration()
-
-        # Initialize blue-green upgrade settings with existing values as defaults
-        blue_green_upgrade_settings = agentpool.blue_green_upgrade_settings
-        if blue_green_upgrade_settings is None:
-            blue_green_upgrade_settings = self.models.AgentPoolBlueGreenUpgradeSettings()
-            
         drain_batch_size = self.context.get_drain_batch_size()
-        if drain_batch_size is not None:
-            blue_green_upgrade_settings.drain_batch_size = drain_batch_size
-            
-        if drain_timeout_bg is not None:
-            blue_green_upgrade_settings.drain_timeout_in_minutes = drain_timeout_bg
-            
-        if batch_soak_duration is not None:
-            blue_green_upgrade_settings.batch_soak_duration_in_minutes = batch_soak_duration
 
-        if final_soak_duration is not None:
-            blue_green_upgrade_settings.final_soak_duration_in_minutes = final_soak_duration            
-            
-        # Set the blue-green upgrade settings as a separate property on agentpool
-        agentpool.upgrade_settings_blue_green = blue_green_upgrade_settings
+        if any([drain_timeout_bg, batch_soak_duration, final_soak_duration, drain_batch_size]):
+            # Initialize blue-green upgrade settings with existing values as defaults
+            upgrade_settings_blue_green = agentpool.upgrade_settings_blue_green
+            if upgrade_settings_blue_green is None:
+                upgrade_settings_blue_green = self.models.AgentPoolBlueGreenUpgradeSettings()
+                
+            if drain_batch_size is not None:
+                upgrade_settings_blue_green.drain_batch_size = drain_batch_size
+                
+            if drain_timeout_bg is not None:
+                upgrade_settings_blue_green.drain_timeout_in_minutes = drain_timeout_bg
+                
+            if batch_soak_duration is not None:
+                upgrade_settings_blue_green.batch_soak_duration_in_minutes = batch_soak_duration
+
+            if final_soak_duration is not None:
+                upgrade_settings_blue_green.final_soak_duration_in_minutes = final_soak_duration            
+                
+            # Set the blue-green upgrade settings as a separate property on agentpool
+            agentpool.upgrade_settings_blue_green = upgrade_settings_blue_green
 
         return agentpool