aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--tests/unit/benchmark/scenarios/networking/test_vnf_generic.py4
-rw-r--r--yardstick/benchmark/scenarios/networking/vnf_generic.py54
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)