Fix DHCPv6 Relay config CLI to update DHCP_RELAY

sonic-net · Dec 13, 2022 · 34aa355 · 34aa355
1 parent 0dc22bd
commit 34aa355
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 10 deletions.
diff --git a/dockers/docker-dhcp-relay/cli-plugin-tests/conftest.py b/dockers/docker-dhcp-relay/cli-plugin-tests/conftest.py
@@ -10,6 +10,11 @@ def mock_cfgdb():
             'Vlan1000': {
                 'dhcp_servers': ['192.0.0.1']
             }
+        },
+        'DHCP_RELAY': {
+            'Vlan1000': {
+                'dhcpv6_servers': ['fc02:2000::100']
+            }
         }
     }
 

diff --git a/dockers/docker-dhcp-relay/cli-plugin-tests/test_config_dhcp_relay.py b/dockers/docker-dhcp-relay/cli-plugin-tests/test_config_dhcp_relay.py
@@ -154,7 +154,7 @@ def test_config_vlan_add_del_dhcpv6_relay_dest(self, mock_cfgdb):
             assert result.exit_code == 0
             assert result.output == config_vlan_add_dhcpv6_relay_output
             assert mock_run_command.call_count == 3
-            db.cfgdb.set_entry.assert_called_once_with('VLAN', 'Vlan1000', {'dhcp_servers': ['192.0.0.1'], 'dhcpv6_servers': ['fc02:2000::1']})
+            db.cfgdb.set_entry.assert_called_once_with('DHCP_RELAY', 'Vlan1000', {'dhcpv6_servers': ['fc02:2000::100','fc02:2000::1']})
 
         db.cfgdb.set_entry.reset_mock()
 
@@ -167,7 +167,7 @@ def test_config_vlan_add_del_dhcpv6_relay_dest(self, mock_cfgdb):
             assert result.exit_code == 0
             assert result.output == config_vlan_del_dhcpv6_relay_output
             assert mock_run_command.call_count == 3
-            db.cfgdb.set_entry.assert_called_once_with('VLAN', 'Vlan1000', {'dhcp_servers': ['192.0.0.1']})
+            db.cfgdb.set_entry.assert_called_once_with('DHCP_RELAY', 'Vlan1000', {'dhcpv6_servers': ['fc02:2000::100']})
 
     def test_config_vlan_add_del_multiple_dhcpv6_relay_dest(self, mock_cfgdb):
         runner = CliRunner()
@@ -183,7 +183,7 @@ def test_config_vlan_add_del_multiple_dhcpv6_relay_dest(self, mock_cfgdb):
             assert result.exit_code == 0
             assert result.output == config_vlan_add_multiple_dhcpv6_relay_output
             assert mock_run_command.call_count == 3
-            db.cfgdb.set_entry.assert_called_once_with('VLAN', 'Vlan1000', {'dhcp_servers': ['192.0.0.1'], 'dhcpv6_servers': ['fc02:2000::1', 'fc02:2000::2', 'fc02:2000::3']})
+            db.cfgdb.set_entry.assert_called_once_with('DHCP_RELAY', 'Vlan1000', {'dhcpv6_servers': ['fc02:2000::100', 'fc02:2000::1', 'fc02:2000::2', 'fc02:2000::3']})
 
         db.cfgdb.set_entry.reset_mock()
 
@@ -196,7 +196,7 @@ def test_config_vlan_add_del_multiple_dhcpv6_relay_dest(self, mock_cfgdb):
             assert result.exit_code == 0
             assert result.output == config_vlan_del_multiple_dhcpv6_relay_output
             assert mock_run_command.call_count == 3
-            db.cfgdb.set_entry.assert_called_once_with('VLAN', 'Vlan1000', {'dhcp_servers': ['192.0.0.1']})
+            db.cfgdb.set_entry.assert_called_once_with('DHCP_RELAY', 'Vlan1000', {'dhcpv6_servers': ['fc02:2000::100']})
 
     def test_config_vlan_remove_nonexist_dhcp_relay_dest(self, mock_cfgdb):
         runner = CliRunner()

diff --git a/dockers/docker-dhcp-relay/cli/config/plugins/dhcp_relay.py b/dockers/docker-dhcp-relay/cli/config/plugins/dhcp_relay.py
@@ -16,15 +16,25 @@ def add_vlan_dhcp_relay_destination(db, vid, dhcp_relay_destination_ips):
     ctx = click.get_current_context()
     added_servers = []
 
+    relay_table = ""
+    dhcpv6_servers = []
+    for relay_addr in dhcp_relay_destination_ips:
+        if clicommon.ipaddress_type(relay_addr) == 4:
+            relay_table = 'VLAN'
+        else:
+            relay_table = 'DHCP_RELAY'
+        break
+
     # Verify vlan is valid
     vlan_name = 'Vlan{}'.format(vid)
-    vlan = db.cfgdb.get_entry('VLAN', vlan_name)
-    if len(vlan) == 0:
+    vlan = db.cfgdb.get_entry(relay_table, vlan_name)
+    if len(vlan) == 0 and relay_table == 'VLAN':
         ctx.fail("{} doesn't exist".format(vlan_name))
 
     # Verify all ip addresses are valid and not exist in DB
     dhcp_servers = vlan.get('dhcp_servers', [])
-    dhcpv6_servers = vlan.get('dhcpv6_servers', [])
+    if len(vlan) != 0:
+        dhcpv6_servers = vlan.get('dhcpv6_servers', [])
 
     for ip_addr in dhcp_relay_destination_ips:
         try:
@@ -46,7 +56,7 @@ def add_vlan_dhcp_relay_destination(db, vid, dhcp_relay_destination_ips):
     if len(dhcpv6_servers):
         vlan['dhcpv6_servers'] = dhcpv6_servers
 
-    db.cfgdb.set_entry('VLAN', vlan_name, vlan)
+    db.cfgdb.set_entry(relay_table, vlan_name, vlan)
 
     if len(added_servers):
         click.echo("Added DHCP relay destination addresses {} to {}".format(added_servers, vlan_name))
@@ -67,9 +77,18 @@ def del_vlan_dhcp_relay_destination(db, vid, dhcp_relay_destination_ips):
 
     ctx = click.get_current_context()
 
+    relay_table = ""
+    dhcpv6_servers = []
+    for relay_addr in dhcp_relay_destination_ips:
+        if clicommon.ipaddress_type(relay_addr) == 4:
+            relay_table = 'VLAN'
+        else:
+            relay_table = 'DHCP_RELAY'
+        break
+
     # Verify vlan is valid
     vlan_name = 'Vlan{}'.format(vid)
-    vlan = db.cfgdb.get_entry('VLAN', vlan_name)
+    vlan = db.cfgdb.get_entry(relay_table, vlan_name)
     if len(vlan) == 0:
         ctx.fail("{} doesn't exist".format(vlan_name))
 
@@ -98,7 +117,7 @@ def del_vlan_dhcp_relay_destination(db, vid, dhcp_relay_destination_ips):
         if 'dhcpv6_servers' in vlan.keys():
             del vlan['dhcpv6_servers']
 
-    db.cfgdb.set_entry('VLAN', vlan_name, vlan)
+    db.cfgdb.set_entry(relay_table, vlan_name, vlan)
     click.echo("Removed DHCP relay destination addresses {} from {}".format(dhcp_relay_destination_ips, vlan_name))
     try:
         click.echo("Restarting DHCP relay service...")