Merge "Fix port-forwarding issue in pod restart tests"
diff --git a/TestON/drivers/common/clidriver.py b/TestON/drivers/common/clidriver.py
index 9c2b0a4..1e4d233 100644
--- a/TestON/drivers/common/clidriver.py
+++ b/TestON/drivers/common/clidriver.py
@@ -1393,12 +1393,15 @@
                         "-n %s" % namespace if namespace else "",
                         podName,
                         portsList )
+
+            self.clearBuffer()
             main.log.info( self.name + ": sending: " + repr( cmdStr ) )
             self.handle.sendline( cmdStr )
             self.handle.expect( "pod/%s" % podName )
             output = self.handle.before + self.handle.after
             i = self.handle.expect( [ "not found", "error", "closed/timedout",
-                                      self.prompt, "The connection to the server", "Forwarding from" ] )
+                                      self.prompt, "The connection to the server",
+                                      "Forwarding from", pexpect.TIMEOUT ] )
             output += self.handle.before + str( self.handle.after )
             # NOTE: This won't clear the buffer entirely, and each time the port forward
             #       is used, another line will be added to the buffer. We need to make
@@ -1409,6 +1412,8 @@
                 self.preDisconnect = self.exitFromProcess
                 self.portForwardList = portsList
                 return main.TRUE
+            elif i == 6:
+                return self.checkPortForward( podName, portsList, kubeconfig, namespace )
             else:
                 main.log.error( self.name + ": Error executing command" )
                 main.log.debug( self.name + ": " + output )
@@ -1442,6 +1447,7 @@
 
         """
         try:
+            main.log.debug( "%s: Checking port-forward session to %s" % ( self.name, podName ) )
             if not portsList:
                 portsList = self.portForwardList
             self.handle.sendline( "" )
@@ -1450,6 +1456,7 @@
             main.log.debug( "%s: %s" % ( self.name, output ) )
             if i == 0:
                 # We are not currently in a port-forwarding session, try to re-establish.
+                main.log.warn( "%s: port-forwarding session to %s closed, attempting to reestablish." % ( self.name, podName ) )
                 return self.kubectlPortForward( podName, portsList, kubeconfig, namespace )
             elif i == 1:
                 # Still in a command, port-forward is probably still active
diff --git a/TestON/tests/USECASE/SegmentRouting/SRStaging/dependencies/SRStagingTest.py b/TestON/tests/USECASE/SegmentRouting/SRStaging/dependencies/SRStagingTest.py
index 5ce8b0d..9cf7f61 100644
--- a/TestON/tests/USECASE/SegmentRouting/SRStaging/dependencies/SRStagingTest.py
+++ b/TestON/tests/USECASE/SegmentRouting/SRStaging/dependencies/SRStagingTest.py
@@ -1128,6 +1128,8 @@
             utilities.assert_equals( expect=main.TRUE, actual=deleted,
                                      onpass="Successfully deleted switch pod",
                                      onfail="Failed to delete switch pod" )
+            # Try to fix ssh connection issue on reconnect
+            controller.k8s.exitFromProcess()
             return k8sNode
         except SkipCase:
             raise