Move the dpu check to eni context

Signed-off-by: Vivek Reddy <[email protected]>

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)

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')