Add with-statment support to VatTerminal.

Change-Id: I7b4b32ce07b9247ccf80bf6b5d3339d00bc0999f
Signed-off-by: Matej Klotton <mklotton@cisco.com>

6 files changed, 77 insertions, 81 deletions

diff --git a/resources/libraries/python/IPv6Setup.py b/resources/libraries/python/IPv6Setup.py
index 86c8876e89..e04668f953 100644
--- a/resources/libraries/python/IPv6Setup.py
+++ b/resources/libraries/python/IPv6Setup.py
@@ -161,15 +161,14 @@ class IPv6Setup(object):
            :type prefix: str
         """
         sw_if_index = Topology.get_interface_sw_index(node, interface)
-        vat = VatTerminal(node)
-        vat.vat_terminal_exec_cmd_from_template('add_ip_address.vat',
-                                                sw_if_index=sw_if_index,
-                                                address=addr,
-                                                prefix_length=prefix)
-        vat.vat_terminal_exec_cmd_from_template('set_if_state.vat',
-                                                sw_if_index=sw_if_index,
-                                                state='admin-up')
-        vat.vat_terminal_close()
+        with VatTerminal(node) as vat:
+            vat.vat_terminal_exec_cmd_from_template('add_ip_address.vat',
+                                                    sw_if_index=sw_if_index,
+                                                    address=addr,
+                                                    prefix_length=prefix)
+            vat.vat_terminal_exec_cmd_from_template('set_if_state.vat',
+                                                    sw_if_index=sw_if_index,
+                                                    state='admin-up')
 
         ssh = SSH()
         ssh.connect(node)
@@ -194,15 +193,14 @@ class IPv6Setup(object):
            :type prefix: str
         """
         sw_if_index = Topology.get_interface_sw_index(node, interface)
-        vat = VatTerminal(node)
-        vat.vat_terminal_exec_cmd_from_template('del_ip_address.vat',
-                                                sw_if_index=sw_if_index,
-                                                address=addr,
-                                                prefix_length=prefix)
-        vat.vat_terminal_exec_cmd_from_template('set_if_state.vat',
-                                                sw_if_index=sw_if_index,
-                                                state='admin-down')
-        vat.vat_terminal_close()
+        with VatTerminal(node) as vat:
+            vat.vat_terminal_exec_cmd_from_template('del_ip_address.vat',
+                                                    sw_if_index=sw_if_index,
+                                                    address=addr,
+                                                    prefix_length=prefix)
+            vat.vat_terminal_exec_cmd_from_template('set_if_state.vat',
+                                                    sw_if_index=sw_if_index,
+                                                    state='admin-down')
 
     @staticmethod
     def vpp_ra_supress_link_layer(node, interface):
diff --git a/resources/libraries/python/InterfaceUtil.py b/resources/libraries/python/InterfaceUtil.py
index c074a88ac3..14473ee4d4 100644
--- a/resources/libraries/python/InterfaceUtil.py
+++ b/resources/libraries/python/InterfaceUtil.py
@@ -66,30 +66,29 @@ class InterfaceUtil(object):
         :raises: RuntimeError if the timeout period value has elapsed.
         """
         if_ready = False
-        vat = VatTerminal(node)
-        not_ready = []
-        start = time()
-        while if_ready != True:
-            out = vat.vat_terminal_exec_cmd('sw_interface_dump')
-            if time() - start > timeout:
+        with VatTerminal(node) as vat:
+            not_ready = []
+            start = time()
+            while not if_ready:
+                out = vat.vat_terminal_exec_cmd('sw_interface_dump')
+                if time() - start > timeout:
+                    for interface in out:
+                        if interface.get('admin_up_down') == 1:
+                            if interface.get('link_up_down') != 1:
+                                logger.debug('{0} link-down'.format(
+                                    interface.get('interface_name')))
+                    raise RuntimeError('timeout, not up {0}'.format(not_ready))
+                not_ready = []
                 for interface in out:
                     if interface.get('admin_up_down') == 1:
                         if interface.get('link_up_down') != 1:
-                            logger.debug('{0} link-down'.format(
-                                interface.get('interface_name')))
-                raise RuntimeError('timeout, not up {0}'.format(not_ready))
-            not_ready = []
-            for interface in out:
-                if interface.get('admin_up_down') == 1:
-                    if interface.get('link_up_down') != 1:
-                        not_ready.append(interface.get('interface_name'))
-            if not not_ready:
-                if_ready = True
-            else:
-                logger.debug('Interfaces still in link-down state: {0}, ' \
-                             'waiting...'.format(not_ready))
-                sleep(1)
-        vat.vat_terminal_close()
+                            not_ready.append(interface.get('interface_name'))
+                if not not_ready:
+                    if_ready = True
+                else:
+                    logger.debug('Interfaces still in link-down state: {0}, '
+                                 'waiting...'.format(not_ready))
+                    sleep(1)
 
     @staticmethod
     def vpp_nodes_interfaces_ready_wait(nodes, timeout=10):
@@ -98,7 +97,7 @@ class InterfaceUtil(object):
 
         :param nodes: List of nodes to wait on.
         :param timeout: Seconds to wait per node for all interfaces to come up.
-        :type node: list
+        :type nodes: list
         :type timeout: int
         :raises: RuntimeError if the timeout period value has elapsed.
         """
@@ -112,7 +111,7 @@ class InterfaceUtil(object):
 
         :param nodes: Nodes in the topology.
         :param timeout: Seconds to wait per node for all interfaces to come up.
-        :type node: dict
+        :type nodes: dict
         :type timeout: int
         :raises: RuntimeError if the timeout period value has elapsed.
         """
diff --git a/resources/libraries/python/L2Util.py b/resources/libraries/python/L2Util.py
index a89130814c..c0a764fa2d 100644
--- a/resources/libraries/python/L2Util.py
+++ b/resources/libraries/python/L2Util.py
@@ -167,11 +167,10 @@ class L2Util(object):
         """
         sw_iface1 = Topology().get_interface_sw_index(node, interface1)
         sw_iface2 = Topology().get_interface_sw_index(node, interface2)
-        vat = VatTerminal(node)
-        vat.vat_terminal_exec_cmd_from_template('l2_xconnect.vat',
-                                                interface1=sw_iface1,
-                                                interface2=sw_iface2)
-        vat.vat_terminal_exec_cmd_from_template('l2_xconnect.vat',
-                                                interface1=sw_iface2,
-                                                interface2=sw_iface1)
-        vat.vat_terminal_close()
+        with VatTerminal(node) as vat:
+            vat.vat_terminal_exec_cmd_from_template('l2_xconnect.vat',
+                                                    interface1=sw_iface1,
+                                                    interface2=sw_iface2)
+            vat.vat_terminal_exec_cmd_from_template('l2_xconnect.vat',
+                                                    interface1=sw_iface2,
+                                                    interface2=sw_iface1)
diff --git a/resources/libraries/python/Routing.py b/resources/libraries/python/Routing.py
index b22516a7c4..7795e3ffcc 100644
--- a/resources/libraries/python/Routing.py
+++ b/resources/libraries/python/Routing.py
@@ -29,17 +29,16 @@ class Routing(object):
         :param prefix_len: Route destination network prefix length.
         :param gateway: Route gateway address.
         :param interface: Route interface.
-        :type node: str
+        :type node: dict
         :type network: str
         :type prefix_len: int
         :type gateway: str
         :type interface: str
         """
         sw_if_index = Topology.get_interface_sw_index(node, interface)
-        vat = VatTerminal(node)
-        vat.vat_terminal_exec_cmd_from_template('add_route.vat',
-                                                network=network,
-                                                prefix_length=prefix_len,
-                                                gateway=gateway,
-                                                sw_if_index=sw_if_index)
-        vat.vat_terminal_close()
+        with VatTerminal(node) as vat:
+            vat.vat_terminal_exec_cmd_from_template('add_route.vat',
+                                                    network=network,
+                                                    prefix_length=prefix_len,
+                                                    gateway=gateway,
+                                                    sw_if_index=sw_if_index)
diff --git a/resources/libraries/python/VatExecutor.py b/resources/libraries/python/VatExecutor.py
index d53e57c309..136b7cc0ef 100644
--- a/resources/libraries/python/VatExecutor.py
+++ b/resources/libraries/python/VatExecutor.py
@@ -34,7 +34,6 @@ def cleanup_vat_json_output(json_output):
 
 
 class VatExecutor(object):
-
     def __init__(self):
         self._stdout = None
         self._stderr = None
@@ -77,11 +76,12 @@ class VatExecutor(object):
         # TODO: download vpp_api_test output file
         # self._delete_files(node, remote_file_path, remote_file_out)
 
-    def execute_script_json_out(self, vat_name, node, timeout=10,):
+    def execute_script_json_out(self, vat_name, node, timeout=10):
         self.execute_script(vat_name, node, timeout, json_out=True)
         self._stdout = cleanup_vat_json_output(self._stdout)
 
-    def _delete_files(self, node, *files):
+    @staticmethod
+    def _delete_files(node, *files):
         ssh = SSH()
         ssh.connect(node)
         files = " ".join([str(x) for x in files])
@@ -111,16 +111,14 @@ class VatExecutor(object):
     def cmd_from_template(node, vat_template_file, **vat_args):
         """Execute VAT script on specified node. This method supports
          script templates with parameters
-        :param node: node in topology on witch the scrtipt is executed
+        :param node: node in topology on witch the script is executed
         :param vat_template_file: template file of VAT script
         :param vat_args: arguments to the template file
         :return: list of json objects returned by VAT
         """
-        vat = VatTerminal(node)
-        ret = vat.vat_terminal_exec_cmd_from_template(vat_template_file,
-                                                      **vat_args)
-        vat.vat_terminal_close()
-        return ret
+        with VatTerminal(node) as vat:
+            return vat.vat_terminal_exec_cmd_from_template(vat_template_file,
+                                                           **vat_args)
 
     @staticmethod
     def copy_config_to_remote(node, local_path, remote_path):
@@ -158,6 +156,12 @@ class VatTerminal(object):
             'sudo -S {vat} json'.format(vat=Constants.VAT_BIN_NAME),
             self.__VAT_PROMPT)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.vat_terminal_close()
+
     def vat_terminal_exec_cmd(self, cmd):
         """Execute command on the opened VAT terminal.
 
@@ -165,11 +169,11 @@ class VatTerminal(object):
 
            :return: Command output in python representation of JSON format.
         """
-        logger.debug("Executing command in VAT terminal: {}".format(cmd));
+        logger.debug("Executing command in VAT terminal: {}".format(cmd))
         out = self._ssh.interactive_terminal_exec_command(self._tty,
                                                           cmd,
                                                           self.__VAT_PROMPT)
-        logger.debug("VAT output: {}".format(out));
+        logger.debug("VAT output: {}".format(out))
         json_out = json.loads(out)
         return json_out
 
diff --git a/resources/libraries/python/VppCounters.py b/resources/libraries/python/VppCounters.py
index b99d65274b..caae06c495 100644
--- a/resources/libraries/python/VppCounters.py
+++ b/resources/libraries/python/VppCounters.py
@@ -83,19 +83,16 @@ class VppCounters(object):
            :type node: dict
            :return: Stats table.
         """
-        vat = VatTerminal(node)
-        vat.vat_terminal_exec_cmd('want_stats enable')
-        for _ in range(0, 12):
-            stats_table = vat.vat_terminal_exec_cmd('dump_stats_table')
-            if_counters = stats_table['interface_counters']
-            if len(if_counters) > 0:
-                self._stats_table = stats_table
-                vat.vat_terminal_close()
-                return stats_table
-            time.sleep(1)
-
-        vat.vat_terminal_close()
-        return None
+        with VatTerminal(node) as vat:
+            vat.vat_terminal_exec_cmd('want_stats enable')
+            for _ in range(0, 12):
+                stats_table = vat.vat_terminal_exec_cmd('dump_stats_table')
+                if_counters = stats_table['interface_counters']
+                if len(if_counters) > 0:
+                    self._stats_table = stats_table
+                    return stats_table
+                time.sleep(1)
+            return None
 
     def vpp_get_ipv4_interface_counter(self, node, interface):
         return self.vpp_get_ipv46_interface_counter(node, interface, False)