update test_update_agentpool_profile_preview.py

Author: lilypan26

src/aks-preview/azext_aks_preview/tests/latest/test_update_agentpool_profile_preview.py

@@ -132,6 +132,7 @@ def test_update_agentpool_profile_preview_default_behavior(self):
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)
         decorator.update_blue_green_upgrade_settings = Mock(return_value=agentpool)
+        decorator.update_gpu_profile = Mock(return_value=agentpool)
 
         # Act
         result = decorator.update_agentpool_profile_preview()
@@ -154,6 +155,7 @@ def test_update_agentpool_profile_preview_default_behavior(self):
         decorator.update_auto_scaler_properties_vms.assert_called_once_with(agentpool)
         decorator.update_upgrade_strategy.assert_called_once_with(agentpool)
         decorator.update_blue_green_upgrade_settings.assert_called_once_with(agentpool)
+        decorator.update_gpu_profile.assert_called_once_with(agentpool)
 
     def test_update_agentpool_profile_preview_with_agentpools_parameter(self):
         """Test update_agentpool_profile_preview with agentpools parameter."""
@@ -194,6 +196,7 @@ def test_update_agentpool_profile_preview_with_agentpools_parameter(self):
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)
         decorator.update_blue_green_upgrade_settings = Mock(return_value=agentpool)
+        decorator.update_gpu_profile = Mock(return_value=agentpool)
 
         # Act
         result = decorator.update_agentpool_profile_preview(agentpools)
@@ -244,6 +247,7 @@ def test_update_agentpool_profile_preview_managed_system_mode(self):
         decorator.update_auto_scaler_properties_vms = Mock()
         decorator.update_upgrade_strategy = Mock()
         decorator.update_blue_green_upgrade_settings = Mock()
+        decorator.update_gpu_profile = Mock()
 
         # Act
         result = decorator.update_agentpool_profile_preview()
@@ -272,6 +276,7 @@ def test_update_agentpool_profile_preview_managed_system_mode(self):
         decorator.update_auto_scaler_properties_vms.assert_not_called()
         decorator.update_upgrade_strategy.assert_not_called()
         decorator.update_blue_green_upgrade_settings.assert_not_called()
+        decorator.update_gpu_profile.assert_not_called()
 
     def test_update_agentpool_profile_preview_managed_system_mode_with_agentpools(self):
         """Test update_agentpool_profile_preview with ManagedSystem mode and agentpools parameter."""
@@ -349,6 +354,7 @@ def test_update_agentpool_profile_preview_system_mode_regular_flow(self):
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)
         decorator.update_blue_green_upgrade_settings = Mock(return_value=agentpool)
+        decorator.update_gpu_profile = Mock(return_value=agentpool)
 
         # Act
         result = decorator.update_agentpool_profile_preview()
@@ -369,6 +375,7 @@ def test_update_agentpool_profile_preview_system_mode_regular_flow(self):
         decorator.update_auto_scaler_properties_vms.assert_called_once_with(agentpool)
         decorator.update_upgrade_strategy.assert_called_once_with(agentpool)
         decorator.update_blue_green_upgrade_settings.assert_called_once_with(agentpool)
+        decorator.update_gpu_profile.assert_called_once_with(agentpool)
 
     def test_update_agentpool_profile_preview_execution_order(self):
         """Test that update methods are called in the correct order."""
@@ -414,6 +421,7 @@ def mock_method(pool):
         decorator.update_auto_scaler_properties_vms = create_mock_update_method("update_auto_scaler_properties_vms")
         decorator.update_upgrade_strategy = create_mock_update_method("update_upgrade_strategy")
         decorator.update_blue_green_upgrade_settings = create_mock_update_method("update_blue_green_upgrade_settings")
+        decorator.update_gpu_profile = create_mock_update_method("update_gpu_profile")
 
         # Act
         decorator.update_agentpool_profile_preview()
@@ -431,6 +439,7 @@ def mock_method(pool):
             "update_auto_scaler_properties_vms",
             "update_upgrade_strategy",
             "update_blue_green_upgrade_settings"
+            "update_gpu_profile",
         ]
         self.assertEqual(call_order, expected_order)
 
@@ -478,6 +487,7 @@ def track_and_return(pool):
         decorator.update_auto_scaler_properties_vms = create_tracking_mock("update_auto_scaler_properties_vms")
         decorator.update_upgrade_strategy = create_tracking_mock("update_upgrade_strategy")
         decorator.update_blue_green_upgrade_settings = create_tracking_mock("update_blue_green_upgrade_settings")        
+        decorator.update_gpu_profile = create_tracking_mock("update_gpu_profile")
 
         # Act
         result = decorator.update_agentpool_profile_preview()