From 977f7a5f5b733428f3db0afbf68502732849b5da Mon Sep 17 00:00:00 2001 From: Vivek Reddy Date: Tue, 6 Aug 2024 20:11:23 +0000 Subject: [PATCH] Move the dpu check to eni context Signed-off-by: Vivek Reddy --- counterpoll/main.py | 17 ++++------------- tests/counterpoll_test.py | 6 +----- 2 files changed, 5 insertions(+), 18 deletions(-) diff --git a/counterpoll/main.py b/counterpoll/main.py index d7b66d5743..be99d10ba9 100644 --- a/counterpoll/main.py +++ b/counterpoll/main.py @@ -19,7 +19,7 @@ def is_dpu(db): """ Check if the device is DPU """ platform_info = device_info.get_platform_info(db) - if platform_info.get('switch_type', '') == 'dpu': + if platform_info.get('switch_type') == 'dpu': return True else: return False @@ -401,16 +401,15 @@ 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() @click.argument('poll_interval', type=click.IntRange(1000, 30000)) @click.pass_context def interval(ctx, poll_interval): """ Set eni counter query interval """ - if not is_dpu(ctx.obj): - click.echo("ENI counters are not supported on non DPU platforms") - exit(1) - eni_info = {} if poll_interval is not None: eni_info['POLL_INTERVAL'] = poll_interval @@ -420,10 +419,6 @@ def interval(ctx, poll_interval): @click.pass_context def enable(ctx): """ Enable eni counter query """ - if not is_dpu(ctx.obj): - click.echo("ENI counters are not supported on non DPU platforms") - exit(1) - eni_info = {} eni_info['FLEX_COUNTER_STATUS'] = 'enable' ctx.obj.mod_entry("FLEX_COUNTER_TABLE", ENI, eni_info) @@ -432,10 +427,6 @@ def enable(ctx): @click.pass_context def disable(ctx): """ Disable eni counter query """ - if not is_dpu(ctx.obj): - click.echo("ENI counters are not supported on non DPU platforms") - exit(1) - eni_info = {} eni_info['FLEX_COUNTER_STATUS'] = 'disable' ctx.obj.mod_entry("FLEX_COUNTER_TABLE", ENI, eni_info) diff --git a/tests/counterpoll_test.py b/tests/counterpoll_test.py index ee7122cbc3..b39f0cf2d1 100644 --- a/tests/counterpoll_test.py +++ b/tests/counterpoll_test.py @@ -65,7 +65,6 @@ def test_show_dpu(self, mock_get_platform_info): mock_get_platform_info.return_value = {'switch_type': 'dpu'} runner = CliRunner() result = runner.invoke(counterpoll.cli.commands["show"], []) - print(result.output) assert result.output == expected_counterpoll_show_dpu def test_port_buffer_drop_interval(self): @@ -250,8 +249,7 @@ def test_update_eni_status(self, status): runner = CliRunner() db = Db() - result = runner.invoke(counterpoll.cli.commands["eni"].commands[status], [], obj=db.cfgdb) - print(result.exit_code, result.output) + 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" @@ -263,7 +261,6 @@ def test_update_eni_status_dpu(self, mock_get_platform_info, status): db = Db() result = runner.invoke(counterpoll.cli.commands["eni"].commands[status], [], obj=db.cfgdb) - print(result.exit_code, result.output) assert result.exit_code == 0 table = db.cfgdb.get_table('FLEX_COUNTER_TABLE') @@ -277,7 +274,6 @@ def test_update_eni_interval(self, mock_get_platform_info): test_interval = "2000" result = runner.invoke(counterpoll.cli.commands["eni"].commands["interval"], [test_interval], obj=db.cfgdb) - print(result.exit_code, result.output) assert result.exit_code == 0 table = db.cfgdb.get_table('FLEX_COUNTER_TABLE')