diff --git a/counterpoll/main.py b/counterpoll/main.py index 2bfcd7377c..6997eaa5d6 100644 --- a/counterpoll/main.py +++ b/counterpoll/main.py @@ -397,15 +397,12 @@ def disable(ctx): # ENI counter commands -@cli.group() +@click.group() @click.pass_context def eni(ctx): """ ENI counter commands """ ctx.obj = ConfigDBConnector() ctx.obj.connect() - if not is_dpu(ctx.obj): - click.echo("ENI counters are not supported on non DPU platforms") - exit(1) @eni.command(name='interval') @@ -534,3 +531,27 @@ def disable(filename): def delay(filename): """ Delay counters in config_db file """ _update_config_db_flex_counter_table("delay", filename) + + +""" +The list of dynamic commands that are added on a specific condition. +Format: + (click group/command, callback function) +""" +dynamic_commands = [ + (eni, is_dpu) +] + + +def register_dynamic_commands(cmds): + """ + Dynamically register commands based on condition callback. + """ + db = ConfigDBConnector() + db.connect() + for cmd, cb in cmds: + if cb(db): + cli.add_command(cmd) + + +register_dynamic_commands(dynamic_commands) diff --git a/tests/counterpoll_test.py b/tests/counterpoll_test.py index 6c165498c5..184c22a137 100644 --- a/tests/counterpoll_test.py +++ b/tests/counterpoll_test.py @@ -4,6 +4,7 @@ import pytest import mock import sys +import importlib from click.testing import CliRunner from shutil import copyfile from utilities_common.db import Db @@ -248,13 +249,15 @@ def test_update_route_counter_interval(self): def test_update_eni_status(self, status): runner = CliRunner() result = runner.invoke(counterpoll.cli, ["eni", status]) - assert result.exit_code == 1 - assert result.output == "ENI counters are not supported on non DPU platforms\n" + assert 'No such command "eni"' in result.output + assert result.exit_code == 2 @pytest.mark.parametrize("status", ["disable", "enable"]) @mock.patch('counterpoll.main.device_info.get_platform_info') def test_update_eni_status_dpu(self, mock_get_platform_info, status): mock_get_platform_info.return_value = {'switch_type': 'dpu'} + importlib.reload(counterpoll) + runner = CliRunner() db = Db() @@ -267,6 +270,8 @@ def test_update_eni_status_dpu(self, mock_get_platform_info, status): @mock.patch('counterpoll.main.device_info.get_platform_info') def test_update_eni_interval(self, mock_get_platform_info): mock_get_platform_info.return_value = {'switch_type': 'dpu'} + importlib.reload(counterpoll) + runner = CliRunner() db = Db() test_interval = "2000"