Merge changes from topics 'YARDSTICK-1286', 'YARDSTICK-1127'

* changes:
  Make "IterationIPC" MQ producer for VNF control messages
  Enable traffic generator PID in "NSPerf" scenario setup
  Enable the MQ producer in "SampleVNFTrafficGen" class
  Add "TrafficGeneratorProducer" for "GenericTrafficGen" instances
  Add "IterationIPC" runner

11 files changed, 410 insertions, 18 deletions

diff --git a/yardstick/tests/functional/common/messaging/test_messaging.py b/yardstick/tests/functional/common/messaging/test_messaging.py
index 99874343b..f3e31e718 100644
--- a/yardstick/tests/functional/common/messaging/test_messaging.py
+++ b/yardstick/tests/functional/common/messaging/test_messaging.py
@@ -32,25 +32,25 @@ class DummyPayload(payloads.Payload):
 class DummyEndpoint(consumer.NotificationHandler):
 
     def info(self, ctxt, **kwargs):
-        if ctxt['pid'] in self._ctx_pids:
-            self._queue.put('ID {}, data: {}, pid: {}'.format(
-                self._id, kwargs['data'], ctxt['pid']))
+        if ctxt['id'] in self._ctx_ids:
+            self._queue.put('Nr {}, data: {}, id: {}'.format(
+                self._id, kwargs['data'], ctxt['id']))
 
 
 class DummyConsumer(consumer.MessagingConsumer):
 
-    def __init__(self, _id, ctx_pids, queue):
+    def __init__(self, _id, ctx_ids, queue):
         self._id = _id
-        endpoints = [DummyEndpoint(_id, ctx_pids, queue)]
-        super(DummyConsumer, self).__init__(TOPIC, ctx_pids, endpoints)
+        endpoints = [DummyEndpoint(_id, ctx_ids, queue)]
+        super(DummyConsumer, self).__init__(TOPIC, ctx_ids, endpoints)
 
 
 class DummyProducer(producer.MessagingProducer):
     pass
 
 
-def _run_consumer(_id, ctx_pids, queue):
-    _consumer = DummyConsumer(_id, ctx_pids, queue)
+def _run_consumer(_id, ctx_ids, queue):
+    _consumer = DummyConsumer(_id, ctx_ids, queue)
     _consumer.start_rpc_server()
     _consumer.wait()
 
@@ -67,8 +67,8 @@ class MessagingTestCase(base.BaseFunctionalTestCase):
         num_consumers = 10
         ctx_1 = 100001
         ctx_2 = 100002
-        producers = [DummyProducer(TOPIC, pid=ctx_1),
-                     DummyProducer(TOPIC, pid=ctx_2)]
+        producers = [DummyProducer(TOPIC, _id=ctx_1),
+                     DummyProducer(TOPIC, _id=ctx_2)]
 
         processes = []
         for i in range(num_consumers):
@@ -91,7 +91,7 @@ class MessagingTestCase(base.BaseFunctionalTestCase):
             output.append(output_queue.get(True, 1))
 
         self.assertEqual(num_consumers * 4, len(output))
-        msg_template = 'ID {}, data: {}, pid: {}'
+        msg_template = 'Nr {}, data: {}, id: {}'
         for i in range(num_consumers):
             for ctx in [ctx_1, ctx_2]:
                 for message in ['message 0', 'message 1']:
diff --git a/yardstick/tests/unit/benchmark/runner/test_base.py b/yardstick/tests/unit/benchmark/runner/test_base.py
index 559c991f3..49ba1efe4 100644
--- a/yardstick/tests/unit/benchmark/runner/test_base.py
+++ b/yardstick/tests/unit/benchmark/runner/test_base.py
@@ -8,12 +8,17 @@
 ##############################################################################
 
 import time
+import uuid
 
 import mock
+from oslo_config import cfg
+import oslo_messaging
 import subprocess
 
 from yardstick.benchmark.runners import base as runner_base
 from yardstick.benchmark.runners import iteration
+from yardstick.common import messaging
+from yardstick.common.messaging import payloads
 from yardstick.tests.unit import base as ut_base
 
 
@@ -94,3 +99,54 @@ class RunnerTestCase(ut_base.BaseUnitTestCase):
 
         with self.assertRaises(NotImplementedError):
             runner._run_benchmark(mock.Mock(), mock.Mock(), mock.Mock(), mock.Mock())
+
+
+class RunnerProducerTestCase(ut_base.BaseUnitTestCase):
+
+    @mock.patch.object(oslo_messaging, 'Target', return_value='rpc_target')
+    @mock.patch.object(oslo_messaging, 'RPCClient')
+    @mock.patch.object(oslo_messaging, 'get_rpc_transport',
+                       return_value='rpc_transport')
+    @mock.patch.object(cfg, 'CONF')
+    def test__init(self, mock_config, mock_transport, mock_rpcclient,
+                   mock_target):
+        _id = uuid.uuid1().int
+        runner_producer = runner_base.RunnerProducer(_id)
+        mock_transport.assert_called_once_with(
+            mock_config, url='rabbit://yardstick:yardstick@localhost:5672/')
+        mock_target.assert_called_once_with(topic=messaging.TOPIC_RUNNER,
+                                            fanout=True,
+                                            server=messaging.SERVER)
+        mock_rpcclient.assert_called_once_with('rpc_transport', 'rpc_target')
+        self.assertEqual(_id, runner_producer._id)
+        self.assertEqual(messaging.TOPIC_RUNNER, runner_producer._topic)
+
+    @mock.patch.object(oslo_messaging, 'Target', return_value='rpc_target')
+    @mock.patch.object(oslo_messaging, 'RPCClient')
+    @mock.patch.object(oslo_messaging, 'get_rpc_transport',
+                       return_value='rpc_transport')
+    @mock.patch.object(payloads, 'RunnerPayload', return_value='runner_pload')
+    def test_start_iteration(self, mock_runner_payload, *args):
+        runner_producer = runner_base.RunnerProducer(uuid.uuid1().int)
+        with mock.patch.object(runner_producer,
+                               'send_message') as mock_message:
+            runner_producer.start_iteration(version=10)
+
+        mock_message.assert_called_once_with(
+            messaging.RUNNER_METHOD_START_ITERATION, 'runner_pload')
+        mock_runner_payload.assert_called_once_with(version=10, data={})
+
+    @mock.patch.object(oslo_messaging, 'Target', return_value='rpc_target')
+    @mock.patch.object(oslo_messaging, 'RPCClient')
+    @mock.patch.object(oslo_messaging, 'get_rpc_transport',
+                       return_value='rpc_transport')
+    @mock.patch.object(payloads, 'RunnerPayload', return_value='runner_pload')
+    def test_stop_iteration(self, mock_runner_payload, *args):
+        runner_producer = runner_base.RunnerProducer(uuid.uuid1().int)
+        with mock.patch.object(runner_producer,
+                               'send_message') as mock_message:
+            runner_producer.stop_iteration(version=15)
+
+        mock_message.assert_called_once_with(
+            messaging.RUNNER_METHOD_STOP_ITERATION, 'runner_pload')
+        mock_runner_payload.assert_called_once_with(version=15, data={})
diff --git a/yardstick/tests/unit/benchmark/runner/test_iteration_ipc.py b/yardstick/tests/unit/benchmark/runner/test_iteration_ipc.py
new file mode 100644
index 000000000..10d14a8a0
--- /dev/null
+++ b/yardstick/tests/unit/benchmark/runner/test_iteration_ipc.py
@@ -0,0 +1,136 @@
+# Copyright (c) 2018 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import multiprocessing
+import time
+import os
+import uuid
+
+import mock
+
+from yardstick.benchmark.runners import iteration_ipc
+from yardstick.common import messaging
+from yardstick.common.messaging import payloads
+from yardstick.tests.unit import base as ut_base
+
+
+class RunnerIterationIPCEndpointTestCase(ut_base.BaseUnitTestCase):
+
+    def setUp(self):
+        self._id = uuid.uuid1().int
+        self._ctx_ids = [uuid.uuid1().int, uuid.uuid1().int]
+        self._queue = multiprocessing.Queue()
+        self.runner = iteration_ipc.RunnerIterationIPCEndpoint(
+            self._id, self._ctx_ids, self._queue)
+        self._kwargs = {'version': 1, 'iteration': 10, 'kpi': {}}
+        self._pload_dict = payloads.TrafficGeneratorPayload.dict_to_obj(
+            self._kwargs).obj_to_dict()
+
+    def test_tg_method_started(self):
+        self._queue.empty()
+        ctxt = {'id': self._ctx_ids[0]}
+        self.runner.tg_method_started(ctxt, **self._kwargs)
+        time.sleep(0.2)
+
+        output = []
+        while not self._queue.empty():
+            output.append(self._queue.get(True, 1))
+
+        self.assertEqual(1, len(output))
+        self.assertEqual(self._ctx_ids[0], output[0]['id'])
+        self.assertEqual(messaging.TG_METHOD_STARTED, output[0]['action'])
+        self.assertEqual(self._pload_dict, output[0]['payload'].obj_to_dict())
+
+    def test_tg_method_finished(self):
+        self._queue.empty()
+        ctxt = {'id': self._ctx_ids[0]}
+        self.runner.tg_method_finished(ctxt, **self._kwargs)
+        time.sleep(0.2)
+
+        output = []
+        while not self._queue.empty():
+            output.append(self._queue.get(True, 1))
+
+        self.assertEqual(1, len(output))
+        self.assertEqual(self._ctx_ids[0], output[0]['id'])
+        self.assertEqual(messaging.TG_METHOD_FINISHED, output[0]['action'])
+        self.assertEqual(self._pload_dict, output[0]['payload'].obj_to_dict())
+
+    def test_tg_method_iteration(self):
+        self._queue.empty()
+        ctxt = {'id': self._ctx_ids[0]}
+        self.runner.tg_method_iteration(ctxt, **self._kwargs)
+        time.sleep(0.2)
+
+        output = []
+        while not self._queue.empty():
+            output.append(self._queue.get(True, 1))
+
+        self.assertEqual(1, len(output))
+        self.assertEqual(self._ctx_ids[0], output[0]['id'])
+        self.assertEqual(messaging.TG_METHOD_ITERATION, output[0]['action'])
+        self.assertEqual(self._pload_dict, output[0]['payload'].obj_to_dict())
+
+
+class RunnerIterationIPCConsumerTestCase(ut_base.BaseUnitTestCase):
+
+    def setUp(self):
+        self._id = uuid.uuid1().int
+        self._ctx_ids = [uuid.uuid1().int, uuid.uuid1().int]
+        self.consumer = iteration_ipc.RunnerIterationIPCConsumer(
+            self._id, self._ctx_ids)
+        self.consumer._queue = mock.Mock()
+
+    def test__init(self):
+        self.assertEqual({self._ctx_ids[0]: [], self._ctx_ids[1]: []},
+                         self.consumer._kpi_per_id)
+
+    def test_is_all_kpis_received_in_iteration(self):
+        payload = payloads.TrafficGeneratorPayload(
+            version=1, iteration=1, kpi={})
+        msg1 = {'action': messaging.TG_METHOD_ITERATION,
+                'id': self._ctx_ids[0], 'payload': payload}
+        msg2 = {'action': messaging.TG_METHOD_ITERATION,
+                'id': self._ctx_ids[1], 'payload': payload}
+        self.consumer.iteration_index = 1
+
+        self.consumer._queue.empty.side_effect = [False, True]
+        self.consumer._queue.get.return_value = msg1
+        self.assertFalse(self.consumer.is_all_kpis_received_in_iteration())
+
+        self.consumer._queue.empty.side_effect = [False, True]
+        self.consumer._queue.get.return_value = msg2
+        self.assertTrue(self.consumer.is_all_kpis_received_in_iteration())
+
+
+class IterationIPCRunnerTestCase(ut_base.BaseUnitTestCase):
+
+    @mock.patch.object(iteration_ipc, '_worker_process')
+    @mock.patch.object(os, 'getpid', return_value=12345678)
+    @mock.patch.object(multiprocessing, 'Process', return_value=mock.Mock())
+    def test__run_benchmark(self, mock_process, mock_getpid, mock_worker):
+        method = 'method'
+        scenario_cfg = {'type': 'scenario_type'}
+        context_cfg = 'context_cfg'
+        name = '%s-%s-%s' % ('IterationIPC', 'scenario_type', 12345678)
+        runner = iteration_ipc.IterationIPCRunner(mock.ANY)
+        mock_getpid.reset_mock()
+
+        runner._run_benchmark('class', method, scenario_cfg, context_cfg)
+        mock_process.assert_called_once_with(
+            name=name,
+            target=mock_worker,
+            args=(runner.result_queue, 'class', method, scenario_cfg,
+                  context_cfg, runner.aborted, runner.output_queue))
+        mock_getpid.assert_called_once()
diff --git a/yardstick/tests/unit/benchmark/scenarios/networking/test_vnf_generic.py b/yardstick/tests/unit/benchmark/scenarios/networking/test_vnf_generic.py
index bb1a7aaca..77a54c0b8 100644
--- a/yardstick/tests/unit/benchmark/scenarios/networking/test_vnf_generic.py
+++ b/yardstick/tests/unit/benchmark/scenarios/networking/test_vnf_generic.py
@@ -553,6 +553,7 @@ class TestNetworkServiceTestCase(unittest.TestCase):
             tgen.verify_traffic = lambda x: verified_dict
             tgen.terminate = mock.Mock(return_value=True)
             tgen.name = "tgen__1"
+            tgen.run_traffic.return_value = 'tg_id'
             vnf = mock.Mock(autospec=GenericVNF)
             vnf.runs_traffic = False
             vnf.terminate = mock.Mock(return_value=True)
@@ -565,7 +566,6 @@ class TestNetworkServiceTestCase(unittest.TestCase):
             self.s.load_vnf_models = mock.Mock(return_value=self.s.vnfs)
             self.s._fill_traffic_profile = \
                 mock.Mock(return_value=TRAFFIC_PROFILE)
-            self.assertIsNone(self.s.setup())
 
     def test_setup_exception(self):
         with mock.patch("yardstick.ssh.SSH") as ssh:
@@ -656,6 +656,9 @@ class TestNetworkServiceTestCase(unittest.TestCase):
         )
         self.assertEqual(self.s.topology, 'fake_nsd')
 
+    def test_get_mq_ids(self):
+        self.assertEqual(self.s._mq_ids, self.s.get_mq_ids())
+
     def test_teardown(self):
         vnf = mock.Mock(autospec=GenericVNF)
         vnf.terminate = mock.Mock(return_value=True)
diff --git a/yardstick/tests/unit/common/messaging/test_payloads.py b/yardstick/tests/unit/common/messaging/test_payloads.py
index 00ec220c9..37b1f1926 100644
--- a/yardstick/tests/unit/common/messaging/test_payloads.py
+++ b/yardstick/tests/unit/common/messaging/test_payloads.py
@@ -44,3 +44,39 @@ class PayloadTestCase(ut_base.BaseUnitTestCase):
         _dict = {'version': 2, 'key1': 'value100', 'key2': 'value200'}
         payload = _DummyPayload.dict_to_obj(_dict)
         self.assertEqual(set(_dict.keys()), payload._fields)
+
+
+class TrafficGeneratorPayloadTestCase(ut_base.BaseUnitTestCase):
+
+    def test_init(self):
+        tg_payload = payloads.TrafficGeneratorPayload(
+            version=1, iteration=10, kpi={'key1': 'value1'})
+        self.assertEqual(1, tg_payload.version)
+        self.assertEqual(10, tg_payload.iteration)
+        self.assertEqual({'key1': 'value1'}, tg_payload.kpi)
+        self.assertEqual(3, len(tg_payload._fields))
+
+    def test__init_missing_required_fields(self):
+        with self.assertRaises(exceptions.PayloadMissingAttributes):
+            payloads.TrafficGeneratorPayload(version=1, iteration=10)
+        with self.assertRaises(exceptions.PayloadMissingAttributes):
+            payloads.TrafficGeneratorPayload(iteration=10, kpi={})
+        with self.assertRaises(exceptions.PayloadMissingAttributes):
+            payloads.TrafficGeneratorPayload(iteration=10)
+
+
+class RunnerPayloadTestCase(ut_base.BaseUnitTestCase):
+
+    def test_init(self):
+        runner_payload = payloads.RunnerPayload(version=5,
+                                                data={'key1': 'value1'})
+        self.assertEqual(5, runner_payload.version)
+        self.assertEqual({'key1': 'value1'}, runner_payload.data)
+
+    def test__init_missing_required_fields(self):
+        with self.assertRaises(exceptions.PayloadMissingAttributes):
+            payloads.RunnerPayload(version=1)
+        with self.assertRaises(exceptions.PayloadMissingAttributes):
+            payloads.RunnerPayload(data=None)
+        with self.assertRaises(exceptions.PayloadMissingAttributes):
+            payloads.RunnerPayload()
diff --git a/yardstick/tests/unit/common/messaging/test_producer.py b/yardstick/tests/unit/common/messaging/test_producer.py
index 0289689dc..22286e5c3 100644
--- a/yardstick/tests/unit/common/messaging/test_producer.py
+++ b/yardstick/tests/unit/common/messaging/test_producer.py
@@ -44,3 +44,10 @@ class MessagingProducerTestCase(ut_base.BaseUnitTestCase):
                 topic='test_topic', fanout=True, server=messaging.SERVER)
             mock_RPCClient.assert_called_once_with('test_rpc_transport',
                                                    'test_Target')
+
+    def test_id(self):
+        with mock.patch.object(oslo_messaging, 'RPCClient'), \
+                mock.patch.object(oslo_messaging, 'get_rpc_transport'), \
+                mock.patch.object(oslo_messaging, 'Target'):
+            msg_producer = _MessagingProducer('topic', 'id_to_check')
+        self.assertEqual('id_to_check', msg_producer.id)
diff --git a/yardstick/tests/unit/network_services/vnf_generic/vnf/test_base.py b/yardstick/tests/unit/network_services/vnf_generic/vnf/test_base.py
index ebedcb451..43e5ac839 100644
--- a/yardstick/tests/unit/network_services/vnf_generic/vnf/test_base.py
+++ b/yardstick/tests/unit/network_services/vnf_generic/vnf/test_base.py
@@ -15,10 +15,15 @@
 
 import multiprocessing
 import os
+import uuid
 
 import mock
+from oslo_config import cfg
+import oslo_messaging
 import unittest
 
+from yardstick.common import messaging
+from yardstick.common.messaging import payloads
 from yardstick.network_services.vnf_generic.vnf import base
 from yardstick.ssh import SSH
 
@@ -140,6 +145,24 @@ VNFD = {
 }
 
 
+class _DummyGenericTrafficGen(base.GenericTrafficGen):  # pragma: no cover
+
+    def run_traffic(self, *args):
+        pass
+
+    def terminate(self):
+        pass
+
+    def collect_kpi(self):
+        pass
+
+    def instantiate(self, *args):
+        pass
+
+    def scale(self, flavor=''):
+        pass
+
+
 class FileAbsPath(object):
     def __init__(self, module_file):
         super(FileAbsPath, self).__init__()
@@ -221,7 +244,7 @@ class TestGenericVNF(unittest.TestCase):
         self.assertEqual(msg, str(exc.exception))
 
 
-class TestGenericTrafficGen(unittest.TestCase):
+class GenericTrafficGenTestCase(unittest.TestCase):
 
     def test_definition(self):
         """Make sure that the abstract class cannot be instantiated"""
@@ -234,3 +257,81 @@ class TestGenericTrafficGen(unittest.TestCase):
                "abstract methods collect_kpi, instantiate, run_traffic, "
                "scale, terminate")
         self.assertEqual(msg, str(exc.exception))
+
+    def test_get_mq_producer_id(self):
+        vnfd = {'benchmark': {'kpi': mock.ANY},
+                'vdu': [{'external-interface': 'ext_int'}]
+                }
+        tg = _DummyGenericTrafficGen('name', vnfd)
+        tg._mq_producer = mock.Mock()
+        tg._mq_producer.get_id.return_value = 'fake_id'
+        self.assertEqual('fake_id', tg.get_mq_producer_id())
+
+
+class TrafficGeneratorProducerTestCase(unittest.TestCase):
+
+    @mock.patch.object(oslo_messaging, 'Target', return_value='rpc_target')
+    @mock.patch.object(oslo_messaging, 'RPCClient')
+    @mock.patch.object(oslo_messaging, 'get_rpc_transport',
+                       return_value='rpc_transport')
+    @mock.patch.object(cfg, 'CONF')
+    def test__init(self, mock_config, mock_transport, mock_rpcclient,
+                   mock_target):
+        _id = uuid.uuid1().int
+        tg_producer = base.TrafficGeneratorProducer(_id)
+        mock_transport.assert_called_once_with(
+            mock_config, url='rabbit://yardstick:yardstick@localhost:5672/')
+        mock_target.assert_called_once_with(topic=messaging.TOPIC_TG,
+                                            fanout=True,
+                                            server=messaging.SERVER)
+        mock_rpcclient.assert_called_once_with('rpc_transport', 'rpc_target')
+        self.assertEqual(_id, tg_producer._id)
+        self.assertEqual(messaging.TOPIC_TG, tg_producer._topic)
+
+    @mock.patch.object(oslo_messaging, 'Target', return_value='rpc_target')
+    @mock.patch.object(oslo_messaging, 'RPCClient')
+    @mock.patch.object(oslo_messaging, 'get_rpc_transport',
+                       return_value='rpc_transport')
+    @mock.patch.object(payloads, 'TrafficGeneratorPayload',
+                       return_value='tg_pload')
+    def test_tg_method_started(self, mock_tg_payload, *args):
+        tg_producer = base.TrafficGeneratorProducer(uuid.uuid1().int)
+        with mock.patch.object(tg_producer, 'send_message') as mock_message:
+            tg_producer.tg_method_started(version=10)
+
+        mock_message.assert_called_once_with(messaging.TG_METHOD_STARTED,
+                                             'tg_pload')
+        mock_tg_payload.assert_called_once_with(version=10, iteration=0,
+                                                kpi={})
+
+    @mock.patch.object(oslo_messaging, 'Target', return_value='rpc_target')
+    @mock.patch.object(oslo_messaging, 'RPCClient')
+    @mock.patch.object(oslo_messaging, 'get_rpc_transport',
+                       return_value='rpc_transport')
+    @mock.patch.object(payloads, 'TrafficGeneratorPayload',
+                       return_value='tg_pload')
+    def test_tg_method_finished(self, mock_tg_payload, *args):
+        tg_producer = base.TrafficGeneratorProducer(uuid.uuid1().int)
+        with mock.patch.object(tg_producer, 'send_message') as mock_message:
+            tg_producer.tg_method_finished(version=20)
+
+        mock_message.assert_called_once_with(messaging.TG_METHOD_FINISHED,
+                                             'tg_pload')
+        mock_tg_payload.assert_called_once_with(version=20, iteration=0,
+                                                kpi={})
+
+    @mock.patch.object(oslo_messaging, 'Target', return_value='rpc_target')
+    @mock.patch.object(oslo_messaging, 'RPCClient')
+    @mock.patch.object(oslo_messaging, 'get_rpc_transport',
+                       return_value='rpc_transport')
+    @mock.patch.object(payloads, 'TrafficGeneratorPayload',
+                       return_value='tg_pload')
+    def test_tg_method_iteration(self, mock_tg_payload, *args):
+        tg_producer = base.TrafficGeneratorProducer(uuid.uuid1().int)
+        with mock.patch.object(tg_producer, 'send_message') as mock_message:
+            tg_producer.tg_method_iteration(100, version=30, kpi={'k': 'v'})
+
+        mock_message.assert_called_once_with(messaging.TG_METHOD_ITERATION,
+                                             'tg_pload')
+        mock_tg_payload.assert_called_once_with(version=30, iteration=100,
+                                                kpi={'k': 'v'})
diff --git a/yardstick/tests/unit/network_services/vnf_generic/vnf/test_sample_vnf.py b/yardstick/tests/unit/network_services/vnf_generic/vnf/test_sample_vnf.py
index 331e80d00..48ae3b505 100644
--- a/yardstick/tests/unit/network_services/vnf_generic/vnf/test_sample_vnf.py
+++ b/yardstick/tests/unit/network_services/vnf_generic/vnf/test_sample_vnf.py
@@ -1090,6 +1090,57 @@ class TestClientResourceHelper(unittest.TestCase):
 
         self.assertIs(client_resource_helper._connect(client), client)
 
+    @mock.patch.object(ClientResourceHelper, '_build_ports')
+    @mock.patch.object(ClientResourceHelper, '_run_traffic_once')
+    def test_run_traffic(self, mock_run_traffic_once, mock_build_ports):
+        client_resource_helper = ClientResourceHelper(mock.Mock())
+        client = mock.Mock()
+        traffic_profile = mock.Mock()
+        mq_producer = mock.Mock()
+        with mock.patch.object(client_resource_helper, '_connect') \
+                as mock_connect, \
+                mock.patch.object(client_resource_helper, '_terminated') \
+                as mock_terminated:
+            mock_connect.return_value = client
+            type(mock_terminated).value = mock.PropertyMock(
+                side_effect=[0, 1, lambda x: x])
+            client_resource_helper.run_traffic(traffic_profile, mq_producer)
+
+        mock_build_ports.assert_called_once()
+        traffic_profile.register_generator.assert_called_once()
+        mq_producer.tg_method_started.assert_called_once()
+        mq_producer.tg_method_finished.assert_called_once()
+        mq_producer.tg_method_iteration.assert_called_once_with(1)
+        mock_run_traffic_once.assert_called_once_with(traffic_profile)
+
+    @mock.patch.object(ClientResourceHelper, '_build_ports')
+    @mock.patch.object(ClientResourceHelper, '_run_traffic_once',
+                       side_effect=Exception)
+    def test_run_traffic_exception(self, mock_run_traffic_once,
+                                   mock_build_ports):
+        client_resource_helper = ClientResourceHelper(mock.Mock())
+        client = mock.Mock()
+        traffic_profile = mock.Mock()
+        mq_producer = mock.Mock()
+        with mock.patch.object(client_resource_helper, '_connect') \
+                as mock_connect, \
+                mock.patch.object(client_resource_helper, '_terminated') \
+                as mock_terminated:
+            mock_connect.return_value = client
+            type(mock_terminated).value = mock.PropertyMock(return_value=0)
+            mq_producer.reset_mock()
+            # NOTE(ralonsoh): "trex_stl_exceptions.STLError" is mocked
+            with self.assertRaises(Exception):
+                client_resource_helper.run_traffic(traffic_profile,
+                                                   mq_producer)
+
+        mock_build_ports.assert_called_once()
+        traffic_profile.register_generator.assert_called_once()
+        mock_run_traffic_once.assert_called_once_with(traffic_profile)
+        mq_producer.tg_method_started.assert_called_once()
+        mq_producer.tg_method_finished.assert_not_called()
+        mq_producer.tg_method_iteration.assert_not_called()
+
 
 class TestRfc2544ResourceHelper(unittest.TestCase):
 
diff --git a/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_prox.py b/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_prox.py
index 050aa4aa0..3e2f598d2 100644
--- a/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_prox.py
+++ b/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_prox.py
@@ -406,7 +406,8 @@ class TestProxTrafficGen(unittest.TestCase):
         sut.setup_helper.prox_config_dict = {}
         sut._connect_client = mock.Mock(autospec=STLClient)
         sut._connect_client.get_stats = mock.Mock(return_value="0")
-        sut._traffic_runner(mock_traffic_profile)
+        sut._setup_mq_producer = mock.Mock(return_value='mq_producer')
+        sut._traffic_runner(mock_traffic_profile, mock.ANY)
 
     @mock.patch('yardstick.network_services.vnf_generic.vnf.prox_helpers.socket')
     @mock.patch(SSH_HELPER)
diff --git a/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_rfc2544_ixia.py b/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_rfc2544_ixia.py
index 42ac40b50..4ade157a3 100644
--- a/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_rfc2544_ixia.py
+++ b/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_rfc2544_ixia.py
@@ -379,7 +379,8 @@ class TestIXIATrafficGen(unittest.TestCase):
                     mock.mock_open(), create=True)
         @mock.patch('yardstick.network_services.vnf_generic.vnf.tg_rfc2544_ixia.LOG.exception')
         def _traffic_runner(*args):
-            result = sut._traffic_runner(mock_traffic_profile)
+            sut._setup_mq_producer = mock.Mock(return_value='mq_producer')
+            result = sut._traffic_runner(mock_traffic_profile, mock.ANY)
             self.assertIsNone(result)
 
         _traffic_runner()
diff --git a/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_trex.py b/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_trex.py
index 350ba8448..700e910f9 100644
--- a/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_trex.py
+++ b/yardstick/tests/unit/network_services/vnf_generic/vnf/test_tg_trex.py
@@ -387,8 +387,9 @@ class TestTrexTrafficGen(unittest.TestCase):
         # must generate cfg before we can run traffic so Trex port mapping is
         # created
         self.sut.resource_helper.generate_cfg()
+        self.sut._setup_mq_producer = mock.Mock()
         with mock.patch.object(self.sut.resource_helper, 'run_traffic'):
-            self.sut._traffic_runner(mock_traffic_profile)
+            self.sut._traffic_runner(mock_traffic_profile, mock.ANY)
 
     def test__generate_trex_cfg(self):
         vnfd = self.VNFD['vnfd:vnfd-catalog']['vnfd'][0]
@@ -453,9 +454,8 @@ class TestTrexTrafficGen(unittest.TestCase):
         self.sut.ssh_helper.run = mock.Mock()
         self.sut._traffic_runner = mock.Mock(return_value=0)
         self.sut.resource_helper.client_started.value = 1
-        result = self.sut.run_traffic(mock_traffic_profile)
+        self.sut.run_traffic(mock_traffic_profile)
         self.sut._traffic_process.terminate()
-        self.assertIsNotNone(result)
 
     def test_terminate(self):
         vnfd = self.VNFD['vnfd:vnfd-catalog']['vnfd'][0]