diff options
| -rw-r--r-- | tests/unit/benchmark/scenarios/networking/test_vnf_generic.py | 4 | ||||
| -rw-r--r-- | yardstick/benchmark/scenarios/networking/vnf_generic.py | 54 |
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) |