diff --git a/config/vlan.py b/config/vlan.py index 671eada8ca..d6332aa64e 100644 --- a/config/vlan.py +++ b/config/vlan.py @@ -1,6 +1,7 @@ import click import ipaddress import utilities_common.cli as clicommon +import utilities_common.dhcp_relay_util as dhcp_relay_util from time import sleep from .utils import log @@ -13,6 +14,11 @@ def vlan(): """VLAN-related configuration tasks""" pass + +def set_dhcp_relay_table(table, config_db, vlan_name, value): + config_db.set_entry(table, vlan_name, value) + + @vlan.command('add') @click.argument('vid', metavar='', required=True, type=int) @clicommon.pass_db @@ -27,8 +33,17 @@ def add_vlan(db, vid): vlan = 'Vlan{}'.format(vid) if clicommon.check_if_vlanid_exist(db.cfgdb, vlan): ctx.fail("{} already exists".format(vlan)) + if clicommon.check_if_vlanid_exist(db.cfgdb, vlan, "DHCP_RELAY"): + ctx.fail("DHCPv6 relay config for {} already exists".format(vlan)) + + # set dhcpv4_relay table + set_dhcp_relay_table('VLAN', db.cfgdb, vlan, {'vlanid': str(vid)}) + + # set dhcpv6_relay table + set_dhcp_relay_table('DHCP_RELAY', db.cfgdb, vlan, {'vlanid': str(vid)}) + # We need to restart dhcp_relay service after dhcpv6_relay config change + dhcp_relay_util.handle_restart_dhcp_relay_service() - db.cfgdb.set_entry('VLAN', vlan, {'vlanid': vid}) @vlan.command('del') @click.argument('vid', metavar='', required=True, type=int) @@ -54,11 +69,18 @@ def del_vlan(db, vid): ctx.fail("{} can not be removed. First remove IP addresses assigned to this VLAN".format(vlan)) keys = [ (k, v) for k, v in db.cfgdb.get_table('VLAN_MEMBER') if k == 'Vlan{}'.format(vid) ] - + if keys: ctx.fail("VLAN ID {} can not be removed. First remove all members assigned to this VLAN.".format(vid)) - - db.cfgdb.set_entry('VLAN', 'Vlan{}'.format(vid), None) + + # set dhcpv4_relay table + set_dhcp_relay_table('VLAN', db.cfgdb, vlan, None) + + # set dhcpv6_relay table + set_dhcp_relay_table('DHCP_RELAY', db.cfgdb, vlan, None) + # We need to restart dhcp_relay service after dhcpv6_relay config change + dhcp_relay_util.handle_restart_dhcp_relay_service() + def restart_ndppd(): verify_swss_running_cmd = "docker container inspect -f '{{.State.Status}}' swss" diff --git a/tests/conftest.py b/tests/conftest.py index 890f8035d0..ad862f1059 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,6 +10,7 @@ from .mock_tables import dbconnector from . import show_ip_route_common import utilities_common.constants as constants +import config.main as config test_path = os.path.dirname(os.path.abspath(__file__)) modules_path = os.path.dirname(test_path) @@ -218,3 +219,14 @@ def setup_ip_route_commands(): import show.main as show return show + + +@pytest.fixture(scope='function') +def mock_restart_dhcp_relay_service(): + print("We are mocking restart dhcp_relay") + origin_func = config.vlan.dhcp_relay_util.handle_restart_dhcp_relay_service + config.vlan.dhcp_relay_util.handle_restart_dhcp_relay_service = mock.MagicMock(return_value=0) + + yield + + config.vlan.dhcp_relay_util.handle_restart_dhcp_relay_service = origin_func diff --git a/tests/vlan_test.py b/tests/vlan_test.py index cfd1f53db4..9539bab827 100644 --- a/tests/vlan_test.py +++ b/tests/vlan_test.py @@ -1,5 +1,6 @@ import os import traceback +import pytest from unittest import mock from click.testing import CliRunner @@ -8,6 +9,18 @@ import show.main as show from utilities_common.db import Db +IP_VERSION_PARAMS_MAP = { + "ipv4": { + "table": "VLAN" + }, + "ipv6": { + "table": "DHCP_RELAY" + } +} +DHCP_RELAY_TABLE_ENTRY = { + "vlanid": "1001" +} + show_vlan_brief_output="""\ +-----------+-----------------+-----------------+----------------+-----------------------+-------------+ | VLAN ID | IP Address | Ports | Port Tagging | DHCP Helper Address | Proxy ARP | @@ -228,6 +241,8 @@ | 4000 | | PortChannel1001 | tagged | | disabled | +-----------+-----------------+-----------------+----------------+-----------------------+-------------+ """ + + class TestVlan(object): @classmethod def setup_class(cls): @@ -392,7 +407,7 @@ def test_config_vlan_add_rif_portchannel_member(self): assert result.exit_code != 0 assert "Error: PortChannel0001 is a router interface!" in result.output - def test_config_vlan_del_vlan(self): + def test_config_vlan_del_vlan(self, mock_restart_dhcp_relay_service): runner = CliRunner() db = Db() obj = {'config_db':db.cfgdb} @@ -450,7 +465,7 @@ def test_config_vlan_del_nonexist_vlan_member(self): assert result.exit_code != 0 assert "Error: Ethernet0 is not a member of Vlan1000" in result.output - def test_config_add_del_vlan_and_vlan_member(self): + def test_config_add_del_vlan_and_vlan_member(self, mock_restart_dhcp_relay_service): runner = CliRunner() db = Db() @@ -493,7 +508,7 @@ def test_config_add_del_vlan_and_vlan_member(self): assert result.exit_code == 0 assert result.output == show_vlan_brief_output - def test_config_add_del_vlan_and_vlan_member_in_alias_mode(self): + def test_config_add_del_vlan_and_vlan_member_in_alias_mode(self, mock_restart_dhcp_relay_service): runner = CliRunner() db = Db() @@ -768,7 +783,7 @@ def test_config_vlan_proxy_arp_with_nonexist_vlan_intf(self): assert result.exit_code != 0 assert "Interface Vlan1001 does not exist" in result.output - def test_config_vlan_proxy_arp_enable(self): + def test_config_vlan_proxy_arp_enable(self, mock_restart_dhcp_relay_service): runner = CliRunner() db = Db() @@ -780,7 +795,7 @@ def test_config_vlan_proxy_arp_enable(self): assert result.exit_code == 0 assert db.cfgdb.get_entry("VLAN_INTERFACE", "Vlan1000") == {"proxy_arp": "enabled"} - def test_config_vlan_proxy_arp_disable(self): + def test_config_vlan_proxy_arp_disable(self, mock_restart_dhcp_relay_service): runner = CliRunner() db = Db() @@ -831,6 +846,39 @@ def test_config_vlan_add_member_of_portchannel(self): assert result.exit_code != 0 assert "Error: Ethernet32 is part of portchannel!" in result.output + @pytest.mark.parametrize("ip_version", ["ipv4", "ipv6"]) + def test_config_add_del_vlan_dhcp_relay(self, ip_version, mock_restart_dhcp_relay_service): + runner = CliRunner() + db = Db() + + # add vlan 1001 + result = runner.invoke(config.config.commands["vlan"].commands["add"], ["1001"], obj=db) + print(result.exit_code) + print(result.output) + assert result.exit_code == 0 + + assert db.cfgdb.get_entry(IP_VERSION_PARAMS_MAP[ip_version]["table"], "Vlan1001") == DHCP_RELAY_TABLE_ENTRY + + # del vlan 1001 + result = runner.invoke(config.config.commands["vlan"].commands["del"], ["1001"], obj=db) + print(result.exit_code) + print(result.output) + + assert "Vlan1001" not in db.cfgdb.get_keys(IP_VERSION_PARAMS_MAP[ip_version]["table"]) + + @pytest.mark.parametrize("ip_version", ["ipv6"]) + def test_config_add_exist_vlan_dhcp_relay(self, ip_version): + runner = CliRunner() + db = Db() + + db.cfgdb.set_entry("DHCP_RELAY", "Vlan1001", {"vlanid": "1001"}) + # add vlan 1001 + result = runner.invoke(config.config.commands["vlan"].commands["add"], ["1001"], obj=db) + print(result.exit_code) + print(result.output) + assert result.exit_code != 0 + assert "DHCPv6 relay config for Vlan1001 already exists" in result.output + @classmethod def teardown_class(cls): os.environ['UTILITIES_UNIT_TESTING'] = "0" diff --git a/utilities_common/cli.py b/utilities_common/cli.py index 50315f321f..6068812a64 100644 --- a/utilities_common/cli.py +++ b/utilities_common/cli.py @@ -245,10 +245,10 @@ def is_vlanid_in_range(vid): return False -def check_if_vlanid_exist(config_db, vlan): +def check_if_vlanid_exist(config_db, vlan, table_name='VLAN'): """Check if vlan id exits in the config db or ot""" - if len(config_db.get_entry('VLAN', vlan)) != 0: + if len(config_db.get_entry(table_name, vlan)) != 0: return True return False diff --git a/utilities_common/dhcp_relay_util.py b/utilities_common/dhcp_relay_util.py new file mode 100644 index 0000000000..b9c0b4e20f --- /dev/null +++ b/utilities_common/dhcp_relay_util.py @@ -0,0 +1,20 @@ +import click +import utilities_common.cli as clicommon + + +def restart_dhcp_relay_service(): + """ + Restart dhcp_relay service + """ + click.echo("Restarting DHCP relay service...") + clicommon.run_command("systemctl stop dhcp_relay", display_cmd=False) + clicommon.run_command("systemctl reset-failed dhcp_relay", display_cmd=False) + clicommon.run_command("systemctl start dhcp_relay", display_cmd=False) + + +def handle_restart_dhcp_relay_service(): + try: + restart_dhcp_relay_service() + except SystemExit as e: + ctx = click.get_current_context() + ctx.fail("Restart service dhcp_relay failed with error {}".format(e))