From 98f2e5c8d6d88155ce70f23ab53ba6f39bc7f9e3 Mon Sep 17 00:00:00 2001
From: Zack Cerza <zack.cerza@inktank.com>
Date: Fri, 19 Sep 2014 09:42:36 -0600
Subject: [PATCH] Raise ConnectionLostError when the socket dies

Signed-off-by: Zack Cerza <zack.cerza@inktank.com>
---
 teuthology/orchestra/run.py           |  6 +++++-
 teuthology/orchestra/test/test_run.py | 22 ++++++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/teuthology/orchestra/run.py b/teuthology/orchestra/run.py
index 2d9edd291c..ea41efda52 100644
--- a/teuthology/orchestra/run.py
+++ b/teuthology/orchestra/run.py
@@ -6,6 +6,7 @@ from paramiko import ChannelFile
 
 import gevent
 import gevent.event
+import socket
 import pipes
 import logging
 import shutil
@@ -309,7 +310,10 @@ def run(
     :param name: Human readable name (probably hostname) of the destination
                  host
     """
-    (host, port) = client.get_transport().getpeername()
+    try:
+        (host, port) = client.get_transport().getpeername()
+    except socket.error:
+        raise ConnectionLostError(command=quote(args))
 
     if name is None:
         name = host
diff --git a/teuthology/orchestra/test/test_run.py b/teuthology/orchestra/test/test_run.py
index 65c08600c4..05de1c4d30 100644
--- a/teuthology/orchestra/test/test_run.py
+++ b/teuthology/orchestra/test/test_run.py
@@ -2,6 +2,7 @@ from cStringIO import StringIO
 
 import fudge
 import logging
+import socket
 
 from .. import run
 from teuthology.exceptions import (CommandCrashedError, CommandFailedError,
@@ -230,6 +231,27 @@ class TestRun(object):
         assert e.command == 'foo'
         assert str(e) == "SSH connection was lost: 'foo'"
 
+    @fudge.with_fakes
+    def test_run_status_lost_socket(self):
+        fudge.clear_expectations()
+        ssh = fudge.Fake('SSHConnection')
+        out = fudge.Fake('ChannelFile').is_a_stub()
+        logger = fudge.Fake('logger').is_a_stub()
+        channel = fudge.Fake('channel')
+        out.has_attr(channel=channel)
+        transport = ssh.expects('get_transport').with_args().returns_fake()
+        transport.expects('getpeername').with_args().raises(socket.error)
+        e = assert_raises(
+            ConnectionLostError,
+            run.run,
+            client=ssh,
+            logger=logger,
+            args=['foo'],
+            )
+
+        assert e.command == 'foo'
+        assert str(e) == "SSH connection was lost: 'foo'"
+
     @fudge.with_fakes
     def test_run_status_lost_nocheck(self):
         fudge.clear_expectations()
-- 
2.39.5