Merge "vnf_generic: convert sshmanager to class"

2 files changed, 29 insertions, 29 deletions

diff --git a/tests/unit/benchmark/scenarios/networking/test_vnf_generic.py b/tests/unit/benchmark/scenarios/networking/test_vnf_generic.py
index f40b8234f..1b02b6eff 100644
--- a/tests/unit/benchmark/scenarios/networking/test_vnf_generic.py
+++ b/tests/unit/benchmark/scenarios/networking/test_vnf_generic.py
@@ -25,7 +25,7 @@ import unittest
 import mock
 
 from yardstick.benchmark.scenarios.networking.vnf_generic import \
-    ssh_manager, NetworkServiceTestCase, IncorrectConfig, IncorrectSetup
+    SshManager, NetworkServiceTestCase, IncorrectConfig, IncorrectSetup
 from yardstick.network_services.collector.subscriber import Collector
 from yardstick.network_services.vnf_generic.vnf.base import \
     GenericTrafficGen, GenericVNF
@@ -318,7 +318,7 @@ class TestNetworkServiceTestCase(unittest.TestCase):
                 mock.Mock(return_value=(0, SYS_CLASS_NET + IP_ADDR_SHOW, ""))
             ssh.return_value = ssh_mock
             for node, node_dict in self.context_cfg["nodes"].items():
-                with ssh_manager(node_dict) as conn:
+                with SshManager(node_dict) as conn:
                     self.assertIsNotNone(conn)
 
     def test___init__(self):
diff --git a/yardstick/benchmark/scenarios/networking/vnf_generic.py b/yardstick/benchmark/scenarios/networking/vnf_generic.py
index d7ba418c3..2956d6d22 100644
--- a/yardstick/benchmark/scenarios/networking/vnf_generic.py
+++ b/yardstick/benchmark/scenarios/networking/vnf_generic.py
@@ -15,7 +15,6 @@
 
 from __future__ import absolute_import
 import logging
-from contextlib import contextmanager
 import yaml
 
 from yardstick.benchmark.scenarios import base
@@ -49,31 +48,32 @@ class IncorrectSetup(Exception):
     pass
 
 
-@contextmanager
-def ssh_manager(node):
-    """
-    args -> network device mappings
-    returns -> ssh connection ready to be used
-    """
-    conn = None
-    try:
-        ssh_port = node.get("ssh_port", ssh.DEFAULT_PORT)
-        conn = ssh.SSH(user=node.get("user", ""),
-                       host=node.get("ip", ""),
-                       password=node.get("password", ""),
-                       port=ssh_port)
-        conn.wait()
-
-    except (SSHError) as error:
-        LOG.info("connect failed to %s, due to %s", node.get("ip", ""), error)
-    try:
-        if conn:
-            yield conn
-        else:
-            yield False
-    finally:
-        if conn:
-            conn.close()
+class SshManager(object):
+    def __init__(self, node):
+        super(SshManager, self).__init__()
+        self.node = node
+        self.conn = None
+
+    def __enter__(self):
+        """
+        args -> network device mappings
+        returns -> ssh connection ready to be used
+        """
+        try:
+            ssh_port = self.node.get("ssh_port", ssh.DEFAULT_PORT)
+            self.conn = ssh.SSH(user=self.node["user"],
+                                host=self.node["ip"],
+                                password=self.node["password"],
+                                port=ssh_port)
+            self.conn.wait()
+        except (SSHError) as error:
+            LOG.info("connect failed to %s, due to %s", self.node["ip"], error)
+        # self.conn defaults to None
+        return self.conn
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.conn:
+            self.conn.close()
 
 
 class NetworkServiceTestCase(base.Scenario):
@@ -208,7 +208,7 @@ class NetworkServiceTestCase(base.Scenario):
         for node, node_dict in context_cfg["nodes"].items():
 
             cmd = "PATH=$PATH:/sbin:/usr/sbin ip addr show"
-            with ssh_manager(node_dict) as conn:
+            with SshManager(node_dict) as conn:
                 exit_status = conn.execute(cmd)[0]
                 if exit_status != 0:
                     raise IncorrectSetup("Node's %s lacks ip tool." % node)