kea: T6211: update dhcp to get pool size for vrf aswell

Author: David Vølker

src/op_mode/dhcp.py

@@ -166,9 +166,13 @@ def _get_formatted_server_leases(raw_data, family='inet'):
     return output
 
 
-def _get_pool_size(pool, family='inet'):
+def _get_pool_size(pool, family='inet', vrf=''):
     v = 'v6' if family == 'inet6' else ''
-    base = f'service dhcp{v}-server shared-network-name {pool}'
+    # if vrf is set get the correct base for config
+    if vrf:
+        base = f'vrf name {vrf} service dhcp{v}-server shared-network-name {pool}'
+    else:
+        base = f'service dhcp{v}-server shared-network-name {pool}'
     size = 0
     subnets = config.list_nodes(f'{base} subnet')
     for subnet in subnets:
@@ -185,7 +189,7 @@ def _get_pool_size(pool, family='inet'):
     return size
 
 
-def _get_raw_server_pool_statistics(config, family='inet', pool=None):
+def _get_raw_server_pool_statistics(config, family='inet', vrf='', pool=None):
     inet_suffix = '6' if family == 'inet6' else '4'
     pools = [pool] if pool else kea_get_dhcp_pools(config, inet_suffix)
 
@@ -269,11 +273,20 @@ def _verify_server(func):
     def _wrapper(*args, **kwargs):
         config = ConfigTreeQuery()
         family = kwargs.get('family')
+        vrf = kwargs.get('vrf')
         v = 'v6' if family == 'inet6' else ''
-        unconf_message = f'DHCP{v} server is not configured'
+
         # Check if config does not exist
-        if not config.exists(f'service dhcp{v}-server'):
-            raise vyos.opmode.UnconfiguredSubsystem(unconf_message)
+        if vrf:
+            unconf_message = f'DHCP{v} server is not configured for VRF {vrf}'
+            if not config.exists(f'vrf name {vrf} service dhcp{v}-server'):
+                raise vyos.opmode.UnconfiguredSubsystem(unconf_message)
+        else:
+            unconf_message = f'DHCP{v} server is not configured'
+            if not config.exists(f'service dhcp{v}-server'):
+                raise vyos.opmode.UnconfiguredSubsystem(unconf_message)
+
+        # return
         return func(*args, **kwargs)
 
     return _wrapper
@@ -321,7 +334,9 @@ def show_server_pool_statistics(
     if pool and active_pools and pool not in active_pools:
         raise vyos.opmode.IncorrectValue(f'DHCP{v} pool "{pool}" does not exist!')
 
-    pool_data = _get_raw_server_pool_statistics(active_config, family=family, pool=pool)
+    pool_data = _get_raw_server_pool_statistics(
+        active_config, family=family, vrf=vrf, pool=pool
+    )
     if raw:
         return pool_data
     else: