[dts] [PATCH 2/2] framework: new function just return dut session message in certain time

[Yong Liu <yong.liu@intel.com>]

From: Marvin Liu <yong.liu@intel.com>

Add new function which will return all the output message before timeout.
This function will be used in scenario like dut start one application, then
tester send packets to dut and need to check the message output from dut.

Optimize basic ssh_pexpect module, make sure buffer has been flushed.
Optimize basic ssh_pexpect module, when session close will raised expection.

Signed-off-by: Marvin Liu <yong.liu@intel.com>

diff --git a/framework/crb.py b/framework/crb.py
index a699bfc..0b63f24 100644
--- a/framework/crb.py
+++ b/framework/crb.py
@@ -69,6 +69,12 @@ class Crb(object):
 
         return self.session.send_expect(cmds, expected, timeout, verify)
 
+    def get_session_output(self, timeout=TIMEOUT):
+        """
+        Get session output message before timeout
+        """
+        return self.session.get_session_before(timeout)
+
     def set_test_types(self, func_tests, perf_tests):
         """
         Enable or disable function/performance test.
diff --git a/framework/ssh_connection.py b/framework/ssh_connection.py
index 18a6517..7286b14 100644
--- a/framework/ssh_connection.py
+++ b/framework/ssh_connection.py
@@ -62,6 +62,11 @@ class SSHConnection(object):
         self.logger.debug(out)
         return out
 
+    def get_session_before(self, timeout=15):
+        out = self.session.get_session_before(timeout)
+        self.logger.debug(out)
+        return out
+
     def close(self):
         self.session.close()
 
diff --git a/framework/ssh_pexpect.py b/framework/ssh_pexpect.py
index 735df44..812d97b 100644
--- a/framework/ssh_pexpect.py
+++ b/framework/ssh_pexpect.py
@@ -2,7 +2,7 @@ import time
 import pexpect
 import pxssh
 from debugger import ignore_keyintr, aware_keyintr
-from exception import TimeoutException, SSHConnectionException
+from exception import TimeoutException, SSHConnectionException, SSHSessionDeadException
 
 """
 Module handle ssh sessions between tester and DUT.
@@ -14,7 +14,7 @@ Aslo support transfer files to tester or DUT.
 class SSHPexpect(object):
 
     def __init__(self, host, username, password):
-        self.magic_prompt = "[MAGIC PROMPT]"
+        self.magic_prompt = "MAGIC PROMPT"
         try:
             self.session = pxssh.pxssh()
             self.username = username
@@ -31,7 +31,7 @@ class SSHPexpect(object):
         self.logger.config_execution(name)
         self.logger.info("ssh %s@%s" % (self.username, self.host))
 
-    def send_expect_base(self, command, expected, timeout=15):
+    def send_expect_base(self, command, expected, timeout):
         ignore_keyintr()
         self.__flush() # clear buffer
         self.session.PROMPT = expected
@@ -54,21 +54,44 @@ class SSHPexpect(object):
         else:
             return ret
 
-    def __flush(self):
+    def get_session_before(self, timeout):
+        """
+        Get all output before timeout
+        """
+        ignore_keyintr()
         self.session.PROMPT = self.magic_prompt
-        self.session.prompt(0.1)
+        try:
+            self.session.prompt(timeout)
+        except Exception as e:
+            pass
+
+        aware_keyintr()
+        before = self.get_output_before()
+        self.__flush()
+        return before
+
+    def __flush(self):
+        """
+        Clear all session buffer
+        """
+        self.session.buffer = ""
+        self.session.before = ""
 
     def __prompt(self, command, timeout):
         if not self.session.prompt(timeout):
             raise TimeoutException(command, self.get_output_all())
 
     def __sendline(self, command):
+        if not self.isalive():
+            raise SSHSessionDeadException(self.host)
         if len(command) == 2 and command.startswith('^'):
             self.session.sendcontrol(command[1])
         else:
             self.session.sendline(command)
 
     def get_output_before(self):
+        if not self.isalive():
+            raise SSHSessionDeadException(self.host)
         self.session.flush()
         before = self.session.before.rsplit('\r\n', 1)
         if before[0] == "[PEXPECT]":
-- 
1.9.3