Added more UTs

Author: Riddhi Nilawar

src/azure-cli/azure/cli/command_modules/network/tests/latest/test_nw_flow_log.py

@@ -220,6 +220,81 @@ def test_nw_flow_log_create_vnetfl_with_record_types(self, resource_group, resou
             self.check('retentionPolicy.enabled', False),
         ])
 
+    
+@ResourceGroupPreparer(name_prefix='test_nw_flow_log_', location='centraluseuap')
+@StorageAccountPreparer(name_prefix='testflowlog', location='centraluseuap', kind='StorageV2')
+def test_nw_flow_log_update_vnetfl_with_record_types(self, resource_group, resource_group_location, storage_account):
+    self.kwargs.update({
+        'rg': resource_group,
+        'location': resource_group_location,
+        'storage_account': storage_account,
+        'vnet': 'vnet1',
+        'subnet': 'subnet1',
+        'nic': 'nic1',
+        'watcher_rg': 'NetworkWatcherRG',
+        'watcher_name': 'NetworkWatcher_{}'.format(resource_group_location),
+        'flow_log': 'flow_log_test_update_rt',
+        'workspace': self.create_random_name('clitest', 20),
+        'record_types': 'B,E'
+    })
+
+    # Prepare target resources
+    self.cmd('network vnet create -g {rg} -n {vnet}')
+    self.cmd('network vnet subnet create -g {rg} --vnet-name {vnet} -n {subnet} --address-prefix 10.0.0.0/24')
+    self.cmd('network nic create -g {rg} -n {nic} --vnet-name {vnet} --subnet {subnet}')
+
+    # Prepare workspace (use eastus for LA workspace)
+    workspace = self.cmd(
+        'monitor log-analytics workspace create '
+        '--resource-group {rg} '
+        '--location eastus '
+        '--workspace-name {workspace} '
+    ).get_output_in_json()
+    self.kwargs.update({'workspace_id': workspace['id']})
+
+    # Create flow log for VNet (no record-types initially)
+    self.cmd(
+        'network watcher flow-log create '
+        '--location {location} '
+        '--resource-group {rg} '
+        '--vnet {vnet} '
+        '--storage-account {storage_account} '
+        '--workspace {workspace_id} '
+        '--name {flow_log} '
+    )
+
+    # Baseline verification
+    res1 = self.cmd(
+        'network watcher flow-log show --location {location} --name {flow_log}'
+    ).get_output_in_json()
+    self.assertEqual(res1['name'], self.kwargs['flow_log'])
+    self.assertEqual(res1['enabled'], True)
+    self.assertTrue(res1['targetResourceId'].endswith(self.kwargs['vnet']))
+
+    # ---- Update: set record types to B,E on the same VNet target ----
+    res2 = self.cmd(
+        'network watcher flow-log update '
+        '--location {location} '
+        '--resource-group {rg} '
+        '--vnet {vnet} '
+        '--name {flow_log} '
+        '--record-types {record_types} '
+    ).get_output_in_json()
+
+    self.assertEqual(res2['name'], self.kwargs['flow_log'])
+    self.assertEqual(res2['enabled'], True)
+    self.assertTrue(res2['targetResourceId'].endswith(self.kwargs['vnet']))
+    self.assertEqual(res2['recordTypes'], self.kwargs['record_types'])
+    # Retention defaults remain unless explicitly changed
+    self.assertEqual(res2['retentionPolicy']['enabled'], False)
+    self.assertEqual(res2['retentionPolicy']['days'], 0)
+
+    # Show again and validate persisted update
+    self.cmd('network watcher flow-log show --location {location} --name {flow_log}', checks=[
+        self.check('name', self.kwargs['flow_log']),
+        self.check('recordTypes', '{record_types}')
+    ])
+
     @ResourceGroupPreparer(name_prefix='test_nw_flow_log_', location='centraluseuap')
     @StorageAccountPreparer(name_prefix='testflowlog', location='centraluseuap', kind='StorageV2')
     def test_nw_flow_log_create_vnetflWithManagedIdentity(self, resource_group, resource_group_location, storage_account):