Move the dpu check to eni context

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

counterpoll/main.py

@@ -16,10 +16,10 @@
 DEFLT_10_SEC= "default (10000)"
 DEFLT_1_SEC = "default (1000)"
 
-def is_dpu(db):
+def is_dpu():
     """ Check if the device is DPU """
-    platform_info = device_info.get_platform_info(db)
-    if platform_info.get('switch_type', '') == 'dpu':
+    platform_info = device_info.get_platform_info()
+    if platform_info.get('switch_type') == 'dpu':
         return True
     else:
         return False
@@ -399,6 +399,9 @@ def disable(ctx):
 @click.pass_context
 def eni(ctx):
     """ ENI counter commands """
+    if not is_dpu():
+        click.echo("ENI counters are not supported on non DPU platforms")
+        exit(1)
     ctx.obj = ConfigDBConnector()
     ctx.obj.connect()
 
@@ -407,10 +410,6 @@ def eni(ctx):
 @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)
@@ -487,7 +478,7 @@ def show():
         data.append(["FLOW_CNT_ROUTE_STAT", route_info.get("POLL_INTERVAL", DEFLT_10_SEC),
                      route_info.get("FLEX_COUNTER_STATUS", DISABLE)])
 
-    if is_dpu(config_db) and eni_info:
+    if is_dpu() and eni_info:
         data.append(["ENI_STAT", eni_info.get("POLL_INTERVAL", DEFLT_1_SEC), eni_info.get("FLEX_COUNTER_STATUS", DISABLE)])
 
     click.echo(tabulate(data, headers=header, tablefmt="simple", missingval=""))

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