Make PFC commands use a class

pfc/main.py

@@ -1,5 +1,4 @@
 #!/usr/bin/env python3
-
 import click
 from swsscommon.swsscommon import ConfigDBConnector
 from tabulate import tabulate
@@ -8,153 +7,167 @@
 ALL_PRIORITIES = [str(x) for x in range(8)]
 PRIORITY_STATUS = ['on', 'off']
 
-def configPfcAsym(interface, pfc_asym):
-    """
-    PFC handler to configure asymmentric PFC.
-    """
-    configdb = ConfigDBConnector()
-    configdb.connect()
 
-    configdb.mod_entry("PORT", interface, {'pfc_asym': pfc_asym})
+class Pfc(object):
+    def __init__(self, cfgdb=None):
+        self.cfgdb = cfgdb
 
+    def configPfcAsym(self, interface, pfc_asym):
+        """
+        PFC handler to configure asymmetric PFC.
+        """
+        configdb = ConfigDBConnector() if self.cfgdb is None else self.cfgdb
+        configdb.connect()
 
-def showPfcAsym(interface):
-    """
-    PFC handler to display asymmetric PFC information.
-    """
-    header = ('Interface', 'Asymmetric')
+        configdb.mod_entry("PORT", interface, {'pfc_asym': pfc_asym})
 
-    configdb = ConfigDBConnector()
-    configdb.connect()
+    def showPfcAsym(self, interface):
+        """
+        PFC handler to display asymmetric PFC information.
+        """
+        header = ('Interface', 'Asymmetric')
 
-    if interface:
-        db_keys = configdb.keys(configdb.CONFIG_DB, 'PORT|{0}'.format(interface))
-    else:
-        db_keys = configdb.keys(configdb.CONFIG_DB, 'PORT|*')
+        configdb = ConfigDBConnector() if self.cfgdb is None else self.cfgdb
+        configdb.connect()
 
-    table = []
-
-    for i in db_keys or [None]:
-        key = None 
-        if i:
-            key = i.split('|')[-1]
+        if interface:
+            db_keys = configdb.keys(configdb.CONFIG_DB, 'PORT|{0}'.format(interface))
+        else:
+            db_keys = configdb.keys(configdb.CONFIG_DB, 'PORT|*')
 
-        if key and key.startswith('Ethernet'):
-            entry = configdb.get_entry('PORT', key)
-            table.append([key, entry.get('pfc_asym', 'N/A')])
+        table = []
 
-    sorted_table = natsorted(table)
+        for i in db_keys or [None]:
+            key = None
+            if i:
+                key = i.split('|')[-1]
 
-    click.echo()
-    click.echo(tabulate(sorted_table, headers=header, tablefmt="simple", missingval=""))
-    click.echo()
+            if key and key.startswith('Ethernet'):
+                entry = configdb.get_entry('PORT', key)
+                table.append([key, entry.get('pfc_asym', 'N/A')])
 
-def configPfcPrio(status, interface, priority):
-    configdb = ConfigDBConnector()
-    configdb.connect()
+        sorted_table = natsorted(table)
 
-    if interface not in configdb.get_keys('PORT_QOS_MAP'):
-        click.echo('Cannot find interface {0}'.format(interface))
-        return 
+        click.echo()
+        click.echo(tabulate(sorted_table, headers=header, tablefmt="simple", missingval=""))
+        click.echo()
 
-    """Current lossless priorities on the interface""" 
-    entry = configdb.get_entry('PORT_QOS_MAP', interface)
-    enable_prio = entry.get('pfc_enable').split(',')
-
-    """Avoid '' in enable_prio"""
-    enable_prio = [x.strip() for x in enable_prio if x.strip()]
-
-    if status == 'on' and priority in enable_prio:
-        click.echo('Priority {0} has already been enabled on {1}'.format(priority, interface))
-        return
-
-    if status == 'off' and priority not in enable_prio:
-        click.echo('Priority {0} is not enabled on {1}'.format(priority, interface))
-        return
-
-    if status == 'on':
-        enable_prio.append(priority)
-
-    else:
-        enable_prio.remove(priority)
-
-    enable_prio.sort()    
-    configdb.mod_entry("PORT_QOS_MAP", interface, {'pfc_enable': ','.join(enable_prio)})
-
-    """Show the latest PFC configuration"""
-    showPfcPrio(interface)
-
-def showPfcPrio(interface):
-    """
-    PFC handler to display PFC enabled priority information.
-    """
-    header = ('Interface', 'Lossless priorities')
-    table = []
+    def configPfcPrio(self, status, interface, priority):
+        configdb = ConfigDBConnector() if self.cfgdb is None else self.cfgdb
+        configdb.connect()
+
+        if interface not in configdb.get_keys('PORT_QOS_MAP'):
+            click.echo('Cannot find interface {0}'.format(interface))
+            return
+
+        """Current lossless priorities on the interface"""
+        entry = configdb.get_entry('PORT_QOS_MAP', interface)
+        enable_prio = entry.get('pfc_enable').split(',')
+
+        """Avoid '' in enable_prio"""
+        enable_prio = [x.strip() for x in enable_prio if x.strip()]
+
+        if status == 'on' and priority in enable_prio:
+            click.echo('Priority {0} has already been enabled on {1}'.format(priority, interface))
+            return
+
+        if status == 'off' and priority not in enable_prio:
+            click.echo('Priority {0} is not enabled on {1}'.format(priority, interface))
+            return
+
+        if status == 'on':
+            enable_prio.append(priority)
+
+        else:
+            enable_prio.remove(priority)
+
+        enable_prio.sort()
+        configdb.mod_entry("PORT_QOS_MAP", interface, {'pfc_enable': ','.join(enable_prio)})
+
+        """Show the latest PFC configuration"""
+        self.showPfcPrio(interface)
 
-    configdb = ConfigDBConnector()
-    configdb.connect()
-
-    """Get all the interfaces with QoS map information"""
-    intfs = configdb.get_keys('PORT_QOS_MAP')
-
-    """The user specifies an interface but we cannot find it"""    
-    if interface and interface not in intfs:
-        click.echo('Cannot find interface {0}'.format(interface))
-        return 
-
-    if interface:
-        intfs = [interface]
-
-    for intf in intfs: 
-        entry = configdb.get_entry('PORT_QOS_MAP', intf)
-        table.append([intf, entry.get('pfc_enable', 'N/A')])
-
-    sorted_table = natsorted(table)
-    click.echo()
-    click.echo(tabulate(sorted_table, headers=header, tablefmt="simple", missingval=""))
-    click.echo()
+    def showPfcPrio(self, interface):
+        """
+        PFC handler to display PFC enabled priority information.
+        """
+        header = ('Interface', 'Lossless priorities')
+        table = []
+
+        configdb = ConfigDBConnector() if self.cfgdb is None else self.cfgdb
+        configdb.connect()
+
+        """Get all the interfaces with QoS map information"""
+        intfs = configdb.get_keys('PORT_QOS_MAP')
+
+        """The user specifies an interface but we cannot find it"""
+        if interface and interface not in intfs:
+            click.echo('Cannot find interface {0}'.format(interface))
+            return
+
+        if interface:
+            intfs = [interface]
+
+        for intf in intfs:
+            entry = configdb.get_entry('PORT_QOS_MAP', intf)
+            table.append([intf, entry.get('pfc_enable', 'N/A')])
+
+        sorted_table = natsorted(table)
+        click.echo()
+        click.echo(tabulate(sorted_table, headers=header, tablefmt="simple", missingval=""))
+        click.echo()
 
 @click.group()
-def cli():
+@click.pass_context
+def cli(ctx):
     """PFC Command Line"""
-    pass
+    # Use the cfgdb object if given as input.
+    cfgdb = None if ctx.obj is None else ctx.obj.cfgdb
+
+    ctx.obj = {'pfc': Pfc(cfgdb)}
 
 @cli.group()
-def config():
+@click.pass_context
+def config(ctx):
     """Config PFC"""
     pass
 
 @cli.group()
-def show():
+@click.pass_context
+def show(ctx):
     """Show PFC information"""
     pass
 
 @click.command()
 @click.argument('status', type=click.Choice(PRIORITY_STATUS))
 @click.argument('interface', type=click.STRING)
-def configAsym(status, interface):
+@click.pass_context
+def configAsym(ctx, status, interface):
     """Configure asymmetric PFC on a given port."""
-    configPfcAsym(interface, status)
+    ctx.obj['pfc'].configPfcAsym(interface, status)
 
 @click.command()
 @click.argument('status', type=click.Choice(PRIORITY_STATUS))
 @click.argument('interface', type=click.STRING)
 @click.argument('priority', type=click.Choice(ALL_PRIORITIES))
-def configPrio(status, interface, priority):
+@click.pass_context
+def configPrio(ctx, status, interface, priority):
     """Configure PFC on a given priority."""
-    configPfcPrio(status, interface, priority)
-    
+    ctx.obj['pfc'].configPfcPrio(status, interface, priority)
+
 @click.command()
 @click.argument('interface', type=click.STRING, required=False)
-def showAsym(interface):
+@click.pass_context
+def showAsym(ctx, interface):
     """Show asymmetric PFC information"""
-    showPfcAsym(interface)
+    ctx.obj['pfc'].showPfcAsym(interface)
 
 @click.command()
 @click.argument('interface', type=click.STRING, required=False)
-def showPrio(interface):
+@click.pass_context
+def showPrio(ctx, interface):
     """Show PFC priority information"""
-    showPfcPrio(interface)
+    ctx.obj['pfc'].showPfcPrio(interface)
 
 config.add_command(configAsym, "asymmetric")
 config.add_command(configPrio, "priority")

tests/pfc_input/assert_show_output.py

@@ -0,0 +1,82 @@
+pfc_asym_cannot_find_intf = """\
+
+Interface    Asymmetric
+-----------  ------------
+
+"""
+
+pfc_cannot_find_intf = """\
+Cannot find interface Ethernet1234
+"""
+
+pfc_show_asymmetric_all = """\
+
+Interface    Asymmetric
+-----------  ------------
+Ethernet0    off
+Ethernet4    off
+Ethernet8    off
+Ethernet12   off
+Ethernet16   off
+Ethernet20   off
+Ethernet24   off
+Ethernet28   off
+Ethernet32   off
+Ethernet36   off
+Ethernet40   off
+Ethernet44   off
+Ethernet48   off
+Ethernet52   off
+Ethernet56   off
+Ethernet60   off
+Ethernet64   off
+Ethernet68   off
+Ethernet72   off
+Ethernet76   off
+Ethernet80   off
+Ethernet84   off
+Ethernet88   off
+Ethernet92   off
+Ethernet96   off
+Ethernet100  off
+Ethernet104  off
+Ethernet108  off
+Ethernet112  off
+Ethernet116  off
+Ethernet120  off
+Ethernet124  off
+
+"""
+
+pfc_show_asymmetric_intf = """\
+
+Interface    Asymmetric
+-----------  ------------
+Ethernet0    off
+
+"""
+
+pfc_show_priority_all = """\
+
+Interface    Lossless priorities
+-----------  ---------------------
+Ethernet0    3,4
+Ethernet4    3,4
+
+"""
+
+pfc_show_priority_intf = """\
+
+Interface    Lossless priorities
+-----------  ---------------------
+Ethernet0    3,4
+
+"""
+
+pfc_config_priority_on = """\
+
+Interface    Lossless priorities
+-----------  ---------------------
+Ethernet0    3,4,5
+
+"""