Awesome working start/stop traffic console

Fixed more stability issues :)
Ready for merging.

3 files changed, 105 insertions, 69 deletions

diff --git a/scripts/automation/trex_control_plane/client/trex_stateless_client.py b/scripts/automation/trex_control_plane/client/trex_stateless_client.py
index 0fff9b36..4e861585 100755
--- a/scripts/automation/trex_control_plane/client/trex_stateless_client.py
+++ b/scripts/automation/trex_control_plane/client/trex_stateless_client.py
@@ -62,8 +62,8 @@ class CTRexStatelessClient(object):
                         continue
                 if bad_ids:
                     # Some port IDs are not according to desires status
-                    raise RuntimeError("The requested method ('{0}') cannot be invoked since port IDs {1} are not "
-                                       "at allowed stated".format(func.__name__, list(bad_ids)))
+                    raise ValueError("The requested method ('{0}') cannot be invoked since port IDs {1} are not "
+                                     "at allowed states".format(func.__name__, list(bad_ids)))
                 else:
                     return func(self, *args, **kwargs)
             return wrapper_f
@@ -127,6 +127,9 @@ class CTRexStatelessClient(object):
     def get_acquired_ports(self):
         return self._conn_handler.keys()
 
+    def get_active_ports(self):
+        return list(self._active_ports)
+
     def acquire(self, port_id, force=False):
         if not self._is_ports_valid(port_id):
             raise ValueError("Provided illegal port id input")
diff --git a/scripts/automation/trex_control_plane/common/text_opts.py b/scripts/automation/trex_control_plane/common/text_opts.py
index 4d1cb0fd..06c2c056 100755
--- a/scripts/automation/trex_control_plane/common/text_opts.py
+++ b/scripts/automation/trex_control_plane/common/text_opts.py
@@ -13,6 +13,8 @@ TEXT_CODES = {'bold': {'start': '\x1b[1m',
                           'end': '\x1b[39m'},
               'green': {'start': '\x1b[32m',
                         'end': '\x1b[39m'},
+              'yellow': {'start': '\x1b[33m',
+                         'end': '\x1b[39m'},
               'underline': {'start': '\x1b[4m',
                             'end': '\x1b[24m'}}
 
@@ -40,6 +42,8 @@ def magenta(text):
 def green(text):
     return text_attribute(text, 'green')
 
+def yellow(text):
+    return text_attribute(text, 'yellow')
 
 def underline(text):
     return text_attribute(text, 'underline')
@@ -54,6 +58,7 @@ def text_attribute(text, attribute):
 FUNC_DICT = {'blue': blue,
              'bold': bold,
              'green': green,
+             'yellow': yellow,
              'cyan': cyan,
              'magenta': magenta,
              'underline': underline,
diff --git a/scripts/automation/trex_control_plane/console/trex_console.py b/scripts/automation/trex_control_plane/console/trex_console.py
index b2eac88e..5c3668c0 100755
--- a/scripts/automation/trex_control_plane/console/trex_console.py
+++ b/scripts/automation/trex_control_plane/console/trex_console.py
@@ -42,41 +42,6 @@ __version__ = "1.0"
 
 LoadedStreamList = namedtuple('LoadedStreamList', ['loaded', 'compiled'])
 
-#
-
-def readch(choices=[]):
-
-    fd = sys.stdin.fileno()
-    old_settings = termios.tcgetattr(fd)
-    try:
-        tty.setraw(sys.stdin.fileno())
-        while True:
-            ch = sys.stdin.read(1)
-            if (ord(ch) == 3) or (ord(ch) == 4):
-                return None
-            if ch in choices:
-                return ch
-    finally:
-        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
-
-    return None
-
-class YesNoMenu(object):
-    def __init__ (self, caption):
-        self.caption = caption
-
-    def show (self):
-        print "{0}".format(self.caption)
-        sys.stdout.write("[Y/y/N/n] : ")
-        ch = readch(choices = ['y', 'Y', 'n', 'N'])
-        if ch == None:
-            return None
-
-        print "\n"
-        if ch == 'y' or ch == 'Y':
-            return True
-        else:
-            return False
 
 class ConfirmMenu(object):
     def __init__ (self, caption):
@@ -280,14 +245,13 @@ class TRexConsole(cmd.Cmd):
             ask = ConfirmMenu('Are you sure you want to acquire all ports ? ')
             rc = ask.show()
             if rc == False:
+                print yellow("[ABORTED]\n")
                 return
             else:
                 port_list = self.stateless_client.get_port_ids()
         else:
             port_list = self.extract_port_ids_from_line(line)
 
-        print "\nTrying to acquire ports: " + (" ".join(str(x) for x in port_list)) + "\n"
-
         # rc, resp_list = self.stateless_client.take_ownership(port_list, force)
         res_ok, log = self.stateless_client.acquire(port_list, force)
         self.prompt_response(log)
@@ -298,11 +262,16 @@ class TRexConsole(cmd.Cmd):
         return
 
 
-    def port_auto_complete(self, text, line, begidx, endidx, acquired=True):
+    def port_auto_complete(self, text, line, begidx, endidx, acquired=True, active=False):
         if acquired:
-            ret_list = [x
-                        for x in map(str, self.stateless_client.get_acquired_ports())
-                        if x.startswith(text)]
+            if not active:
+                ret_list = [x
+                            for x in map(str, self.stateless_client.get_acquired_ports())
+                            if x.startswith(text)]
+            else:
+                ret_list = [x
+                            for x in map(str, self.stateless_client.get_active_ports())
+                            if x.startswith(text)]
         else:
             ret_list = [x
                         for x in map(str, self.stateless_client.get_port_ids())
@@ -328,19 +297,24 @@ class TRexConsole(cmd.Cmd):
             ask = ConfirmMenu('Are you sure you want to release all acquired ports? ')
             rc = ask.show()
             if rc == False:
+                print yellow("[ABORTED]\n")
                 return
             else:
                 port_list = self.stateless_client.get_acquired_ports()
         else:
             port_list = self.extract_port_ids_from_line(line)
 
-        res_ok, log = self.stateless_client.release(port_list)
-        self.prompt_response(log)
-        if not res_ok:
+        try:
+            res_ok, log = self.stateless_client.release(port_list)
+            self.prompt_response(log)
+            if not res_ok:
+                print format_text("[FAILED]\n", 'red', 'bold')
+                return
+            print format_text("[SUCCESS]\n", 'green', 'bold')
+        except ValueError as e:
+            print magenta(str(e))
             print format_text("[FAILED]\n", 'red', 'bold')
             return
-        print format_text("[SUCCESS]\n", 'green', 'bold')
-        return
 
     def complete_release(self, text, line, begidx, endidx):
         return self.port_auto_complete(text, line, begidx, endidx)
@@ -489,7 +463,7 @@ class TRexConsole(cmd.Cmd):
     def do_stream_db_add(self, line):
         '''Loads a YAML stream list serialization into user console \n'''
         args = line.split()
-        if args >= 2:
+        if len(args) >= 2:
             name = args[0]
             yaml_path = args[1]
             try:
@@ -512,22 +486,10 @@ class TRexConsole(cmd.Cmd):
                 print "adding new stream failed due to the following error:\n", str(e)
                 print format_text("[FAILED]\n", 'red', 'bold')
 
-            # if name in self.user_streams:
-            #     print "Picked name already exist. Please pick another name."
-            # else:
-            #     try:
-            #         compiled_streams = stream_list.compile_streams()
-            #         self.user_streams[name] = LoadedStreamList(loaded_obj,
-            #                                                    [StreamPack(v.stream_id, v.stream.dump())
-            #                                                     for k, v in compiled_streams.items()])
-            #
-            #         print "Stream list '{0}' loaded successfully".format(name)
-            #     except Exception as e:
-            #         raise
             return
         else:
-            print "please provide load name and YAML path, separated by space.\n" \
-                  "Optionally, you may provide a third argument to specify multiplier."
+            print magenta("please provide load name and YAML path, separated by space.\n"
+                          "Optionally, you may provide a third argument to specify multiplier.\n")
 
     @staticmethod
     def tree_autocomplete(text):
@@ -573,7 +535,7 @@ class TRexConsole(cmd.Cmd):
     def do_stream_db_remove(self, line):
         args = line.split()
         if args:
-            removed_streams = self.streams_db.remove_stream_packs(args)
+            removed_streams = self.streams_db.remove_stream_packs(*args)
             if removed_streams:
                 print green("The following stream packs were removed:")
                 print bold(", ".join(sorted(removed_streams)))
@@ -604,9 +566,10 @@ class TRexConsole(cmd.Cmd):
                 print format_text("[FAILED]\n", 'red', 'bold')
                 return
             if args[0] == "all":
-                ask = YesNoMenu('Are you sure you want to release all acquired ports ? ')
+                ask = ConfirmMenu('Are you sure you want to release all acquired ports? ')
                 rc = ask.show()
                 if rc == False:
+                    print yellow("[ABORTED]\n")
                     return
                 else:
                     port_list = self.stateless_client.get_acquired_ports()
@@ -631,13 +594,15 @@ class TRexConsole(cmd.Cmd):
                                                           miss=list(set(port_list).difference(owned)))
                 print format_text("[FAILED]\n", 'red', 'bold')
         else:
-            print "Please provide list name and ports to attach to, " \
-                  "or specify 'all' to attach all owned ports."
+            print magenta("Please provide list name and ports to attach to, "
+                          "or specify 'all' to attach all owned ports.\n")
 
     def complete_attach(self, text, line, begidx, endidx):
         arg_num = len(line.split()) - 1
         if arg_num == 1:
             # return optional streams packs
+            if line.endswith(" "):
+                return self.port_auto_complete(text, line, begidx, endidx)
             return [x
                     for x in self.streams_db.get_loaded_streams_names()
                     if x.startswith(text)]
@@ -675,14 +640,13 @@ class TRexConsole(cmd.Cmd):
             ask = ConfirmMenu('Are you sure you want to remove all stream packs from all acquired ports? ')
             rc = ask.show()
             if rc == False:
+                print yellow("[ABORTED]\n")
                 return
             else:
                 port_list = self.stateless_client.get_acquired_ports()
         else:
             port_list = self.extract_port_ids_from_line(line)
 
-        print "\nTrying to remove all streams at ports: " + (" ".join(str(x) for x in port_list)) + "\n"
-
         # rc, resp_list = self.stateless_client.take_ownership(port_list, force)
         res_ok, log = self.stateless_client.remove_all_streams(port_list)
         self.prompt_response(log)
@@ -695,6 +659,70 @@ class TRexConsole(cmd.Cmd):
     def complete_remove_all_streams(self, text, line, begidx, endidx):
         return self.port_auto_complete(text, line, begidx, endidx)
 
+    def do_start_traffic(self, line):
+        # make sure that the user wants to acquire all
+        args = line.split()
+        if len(args) < 1:
+            print magenta("Please provide a list of ports separated by spaces, "
+                          "or specify 'all' to start traffic on all acquired ports")
+            return
+        if args[0] == "all":
+            ask = ConfirmMenu('Are you sure you want to start traffic at all acquired ports? ')
+            rc = ask.show()
+            if rc == False:
+                print yellow("[ABORTED]\n")
+                return
+            else:
+                port_list = self.stateless_client.get_acquired_ports()
+        else:
+            port_list = self.extract_port_ids_from_line(line)
+
+        try:
+            res_ok, log = self.stateless_client.start_traffic(port_list)
+            self.prompt_response(log)
+            if not res_ok:
+                print format_text("[FAILED]\n", 'red', 'bold')
+                return
+            print format_text("[SUCCESS]\n", 'green', 'bold')
+        except ValueError as e:
+            print magenta(str(e))
+            print format_text("[FAILED]\n", 'red', 'bold')
+
+    def complete_start_traffic(self, text, line, begidx, endidx):
+        return self.port_auto_complete(text, line, begidx, endidx)
+
+    def do_stop_traffic(self, line):
+        # make sure that the user wants to acquire all
+        args = line.split()
+        if len(args) < 1:
+            print magenta("Please provide a list of ports separated by spaces, "
+                          "or specify 'all' to stop traffic on all acquired ports")
+            return
+        if args[0] == "all":
+            ask = ConfirmMenu('Are you sure you want to start traffic at all acquired ports? ')
+            rc = ask.show()
+            if rc == False:
+                print yellow("[ABORTED]\n")
+                return
+            else:
+                port_list = self.stateless_client.get_active_ports()
+        else:
+            port_list = self.extract_port_ids_from_line(line)
+
+        try:
+            res_ok, log = self.stateless_client.stop_traffic(port_list)
+            self.prompt_response(log)
+            if not res_ok:
+                print format_text("[FAILED]\n", 'red', 'bold')
+                return
+            print format_text("[SUCCESS]\n", 'green', 'bold')
+        except ValueError as e:
+            print magenta(str(e))
+            print format_text("[FAILED]\n", 'red', 'bold')
+
+    def complete_stop_traffic(self, text, line, begidx, endidx):
+        return self.port_auto_complete(text, line, begidx, endidx, active=True)
+
     # adds a very simple stream
     def do_add_simple_stream(self, line):
         if line == "":