update agentpool update decorator and add tests

Author: lilypan26

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -1644,6 +1644,16 @@ def update_network_profile(self, agentpool: AgentPool) -> AgentPool:
             agentpool.network_profile.allowed_host_ports = allowed_host_ports
         return agentpool
 
+    def update_gpu_profile(self, agentpool: AgentPool) -> AgentPool:
+        self._ensure_agentpool(agentpool)
+
+        gpu_driver = self.context.get_gpu_driver()
+        if gpu_driver is not None:
+            if agentpool.gpu_profile is None:
+                agentpool.gpu_profile = self.models.GPUProfile()
+            agentpool.gpu_profile.driver = gpu_driver
+        return agentpool
+
     def update_artifact_streaming(self, agentpool: AgentPool) -> AgentPool:
         """Update artifact streaming property for the AgentPool object.
         :return: the AgentPool object
@@ -1807,6 +1817,9 @@ def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -
         # update blue-green upgrade settings
         agentpool = self.update_blue_green_upgrade_settings(agentpool)
 
+        # update gpu profile
+        agentpool = self.update_gpu_profile(agentpool)
+
         return agentpool
 
     def update_auto_scaler_properties(self, agentpool: AgentPool) -> AgentPool:

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

@@ -2660,6 +2660,31 @@ def common_update_localdns_profile(self):
         finally:
             os.unlink(config_file_path)
 
+    def common_update_gpu_profile(self):
+        dec_1 = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"gpu_driver": "None"},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        # fail on passing the wrong agentpool object
+        with self.assertRaises(CLIInternalError):
+            dec_1.set_up_gpu_profile(None)
+        agentpool_1 = self.create_initialized_agentpool_instance(
+            gpu_profile=self.models.GPUProfile(
+                driver="Install",
+            )
+        )
+        dec_1.context.attach_agentpool(agentpool_1)
+        dec_agentpool_1 = dec_1.update_gpu_profile(agentpool_1)
+        ground_truth_agentpool_1 = self.create_initialized_agentpool_instance(
+            gpu_profile=self.models.GPUProfile(
+                driver="None",
+            )
+        )
+        self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1)        
+
     def common_test_process_dns_overrides_helper(self):
         from azext_aks_preview._helpers import process_dns_overrides
 
@@ -2744,6 +2769,9 @@ def test_update_blue_green_upgrade_settings(self):
     def test_update_localdns_profile(self):
         self.common_update_localdns_profile()
 
+    def test_update_gpu_profile(self):
+        self.common_update_gpu_profile()
+
     def test_process_dns_overrides_helper(self):
         self.common_test_process_dns_overrides_helper()