From 11bcf4ca8fed5259e321c535bf90d0442e9b9746 Mon Sep 17 00:00:00 2001
From: imarom <imarom@cisco.com>
Date: Mon, 1 Feb 2016 08:59:38 -0500
Subject: fix for http://trex-tgn.cisco.com/youtrack/issue/trex-174

---
 .../client/trex_stateless_client.py                | 35 ++++++++++++++--------
 .../client/trex_stateless_sim.py                   |  8 ++---
 2 files changed, 26 insertions(+), 17 deletions(-)

(limited to 'scripts/automation/trex_control_plane/client')

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 78adbcc2..506decfe 100755
--- a/scripts/automation/trex_control_plane/client/trex_stateless_client.py
+++ b/scripts/automation/trex_control_plane/client/trex_stateless_client.py
@@ -25,6 +25,7 @@ from trex_port import Port
 from common.trex_types import *
 from common.trex_stl_exceptions import *
 from trex_async_client import CTRexAsyncClient
+from yaml import YAMLError
 
 
 ############################     logger     #############################
@@ -436,9 +437,6 @@ class STLClient(object):
         self.stats_generator = trex_stats.CTRexInfoGenerator(self.global_stats,
                                                               self.ports)
 
-        # stream DB
-        self.streams_db = CStreamsDB()
-
  
  
     ############# private functions - used by the class itself ###########
@@ -1301,17 +1299,28 @@ class STLClient(object):
         if not rc:
             raise STLArgumentError('ports', ports, valid_values = self.get_all_ports())
 
-        # load the YAML
+
+        streams = None
+
+        # try YAML
         try:
-            streams_pack = self.streams_db.load_yaml_file(filename)
-        except Exception as e:
-            raise STLError(str(e))
-
-        # HACK - convert the stream pack to simple streams
-        streams = []
-        for stream in streams_pack.compiled:
-            s = HACKSTLStream(stream)
-            streams.append(s)
+            streams_db = CStreamsDB()
+            stream_list = streams_db.load_yaml_file(filename)
+
+            # convert to new style stream object
+            streams = [HACKSTLStream(stream) for stream in stream_list.compiled]
+        except YAMLError:
+            # try python
+            try:
+                basedir = os.path.dirname(filename)
+                sys.path.append(basedir)
+                file    = os.path.basename(filename).split('.')[0]
+                module = __import__(file, globals(), locals(), [], -1)
+
+                streams = module.register().get_streams()
+
+            except (AttributeError, ImportError):
+                raise STLError("bad format input file '{0}'".format(filename))
 
         self.add_streams(streams, ports)
 
diff --git a/scripts/automation/trex_control_plane/client/trex_stateless_sim.py b/scripts/automation/trex_control_plane/client/trex_stateless_sim.py
index 7f65996d..d8f6ed92 100644
--- a/scripts/automation/trex_control_plane/client/trex_stateless_sim.py
+++ b/scripts/automation/trex_control_plane/client/trex_stateless_sim.py
@@ -25,7 +25,7 @@ except ImportError:
     import client.outer_packages
 
 from common.trex_stl_exceptions import STLError
-from yaml.scanner import ScannerError
+from yaml import YAMLError
 from common.trex_streams import *
 from client_utils import parsing_opts
 
@@ -103,7 +103,7 @@ class STLSim(object):
 
             # convert to new style stream object
             return [HACKSTLStream(stream) for stream in stream_list.compiled]
-        except ScannerError:
+        except YAMLError:
             pass
 
         # try python
@@ -111,12 +111,12 @@ class STLSim(object):
             basedir = os.path.dirname(input_file)
             sys.path.append(basedir)
 
-            file    = os.path.basename(input_file).split('.py')[0]
+            file    = os.path.basename(input_file).split('.')[0]
             module = __import__(file, globals(), locals(), [], -1)
 
             return module.register().get_streams()
 
-        except AttributeError:
+        except (AttributeError, ImportError):
             pass
 
         raise STLError("bad format input file '{0}'".format(input_file))
-- 
cgit 1.2.3-korg