CSIT-866: wrk onboarding in CSIT

[csit.git] / resources / libraries / python / TrafficGenerator.py

diff --git a/resources/libraries/python/TrafficGenerator.py b/resources/libraries/python/TrafficGenerator.py
index a808338..f363fe3 100644 (file)
--- a/resources/libraries/python/TrafficGenerator.py
+++ b/resources/libraries/python/TrafficGenerator.py
@@ -185,8 +185,9 @@ class TrafficGenerator(object):
             ssh.connect(tg_node)
 
             (ret, _, _) = ssh.exec_command(
-                "sudo -E sh -c '{}/resources/tools/trex/"
-                "trex_installer.sh'".format(Constants.REMOTE_FW_DIR),
+                "sudo -E sh -c '{0}/resources/tools/trex/"
+                "trex_installer.sh {1}'".format(Constants.REMOTE_FW_DIR,
+                                                Constants.TREX_INSTALL_VERSION),
                 timeout=1800)
             if int(ret) != 0:
                 raise RuntimeError('TRex installation failed.')
@@ -306,6 +307,24 @@ class TrafficGenerator(object):
             # critical error occurred
             raise RuntimeError('t-rex-64 startup failed')
 
+    @staticmethod
+    def is_trex_running(node):
+        """Check if TRex is running using pidof.
+
+        :param node: Traffic generator node.
+        :type node: dict
+        :returns: True if TRex is running otherwise False.
+        :rtype: bool
+        :raises: RuntimeError if node type is not a TG.
+        """
+        if node['type'] != NodeType.TG:
+            raise RuntimeError('Node type is not a TG')
+
+        ssh = SSH()
+        ssh.connect(node)
+        ret, _, _ = ssh.exec_command_sudo("pidof t-rex")
+        return bool(int(ret) == 0)
+
     @staticmethod
     def teardown_traffic_generator(node):
         """TG teardown.