From 384e1a92b716250c168f5dedc1f9693111f81423 Mon Sep 17 00:00:00 2001
From: Lance Stout <lance@dingus.local>
Date: Fri, 27 May 2011 11:01:30 -0700
Subject: [PATCH] Added support for testind disconnect errors.

---
 sleekxmpp/test/livesocket.py     | 12 ++++++++++++
 sleekxmpp/test/mocksocket.py     | 14 ++++++++++++++
 sleekxmpp/test/sleektest.py      |  7 +++++++
 sleekxmpp/xmlstream/xmlstream.py |  3 ++-
 tests/test_stream.py             | 21 ++++++++++++++++++++-
 5 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/sleekxmpp/test/livesocket.py b/sleekxmpp/test/livesocket.py
index 3e0f213..7dd4693 100644
--- a/sleekxmpp/test/livesocket.py
+++ b/sleekxmpp/test/livesocket.py
@@ -58,6 +58,18 @@ class TestLiveSocket(object):
     # ------------------------------------------------------------------
     # Testing Interface
 
+    def disconnect_errror(self):
+        """
+        Used to simulate a socket disconnection error.
+
+        Not used by live sockets.
+        """
+        try:
+            self.socket.shutdown()
+            self.socket.close()
+        except:
+            pass
+
     def next_sent(self, timeout=None):
         """
         Get the next stanza that has been sent.
diff --git a/sleekxmpp/test/mocksocket.py b/sleekxmpp/test/mocksocket.py
index e3ddd70..a2af8d6 100644
--- a/sleekxmpp/test/mocksocket.py
+++ b/sleekxmpp/test/mocksocket.py
@@ -39,6 +39,7 @@ class TestSocket(object):
         self.recv_queue = queue.Queue()
         self.send_queue = queue.Queue()
         self.is_live = False
+        self.disconnected = False
 
     def __getattr__(self, name):
         """
@@ -89,6 +90,13 @@ class TestSocket(object):
         """
         self.recv_queue.put(data)
 
+    def disconnect_error(self):
+        """
+        Simulate a disconnect error by raising a socket.error exception
+        for any current or further socket operations.
+        """
+        self.disconnected = True
+
     # ------------------------------------------------------------------
     # Socket Interface
 
@@ -99,6 +107,8 @@ class TestSocket(object):
         Arguments:
             Placeholders. Same as for socket.Socket.recv.
         """
+        if self.disconnected:
+            raise socket.error
         return self.read(block=True)
 
     def send(self, data):
@@ -108,6 +118,8 @@ class TestSocket(object):
         Arguments:
             data -- String value to write.
         """
+        if self.disconnected:
+            raise socket.error
         self.send_queue.put(data)
 
     # ------------------------------------------------------------------
@@ -132,6 +144,8 @@ class TestSocket(object):
             timeout -- Time in seconds a block should last before
                        returning None.
         """
+        if self.disconnected:
+            raise socket.error
         if timeout is not None:
             block = True
         try:
diff --git a/sleekxmpp/test/sleektest.py b/sleekxmpp/test/sleektest.py
index 24af1e7..8cf7b70 100644
--- a/sleekxmpp/test/sleektest.py
+++ b/sleekxmpp/test/sleektest.py
@@ -259,6 +259,13 @@ class SleekTest(unittest.TestCase):
     # ------------------------------------------------------------------
     # Methods for simulating stanza streams.
 
+    def stream_disconnect(self):
+        """
+        Simulate a stream disconnection.
+        """
+        if self.xmpp:
+            self.xmpp.socket.disconnect_error()
+
     def stream_start(self, mode='client', skip=True, header=None,
                            socket='mock', jid='tester@localhost',
                            password='test', server='localhost',
diff --git a/sleekxmpp/xmlstream/xmlstream.py b/sleekxmpp/xmlstream/xmlstream.py
index 1c16556..468db03 100644
--- a/sleekxmpp/xmlstream/xmlstream.py
+++ b/sleekxmpp/xmlstream/xmlstream.py
@@ -978,7 +978,8 @@ class XMLStream(object):
                 log.debug("SEND: %s" % data)
                 try:
                     self.socket.send(data.encode('utf-8'))
-                except:
+                except Socket.error as serr:
+                    self.event('socket_error', serr)
                     log.warning("Failed to send %s" % data)
                     self.disconnect(self.auto_reconnect)
         except KeyboardInterrupt:
diff --git a/tests/test_stream.py b/tests/test_stream.py
index f91f71f..deac24a 100644
--- a/tests/test_stream.py
+++ b/tests/test_stream.py
@@ -1,5 +1,5 @@
+import time
 from sleekxmpp.test import *
-import sleekxmpp.plugins.xep_0033 as xep_0033
 
 
 class TestStreamTester(SleekTest):
@@ -57,4 +57,23 @@ class TestStreamTester(SleekTest):
         self.stream_start(mode='client', skip=False)
         self.send_header(sto='localhost')
 
+    def testStreamDisconnect(self):
+        """Test that the test socket can simulate disconnections."""
+        self.stream_start()
+        events = set()
+
+        def stream_error(event):
+            events.add('socket_error')
+
+        self.xmpp.add_event_handler('socket_error', stream_error)
+
+        self.stream_disconnect()
+        self.xmpp.send_raw('  ')
+
+        time.sleep(.1)
+
+        self.failUnless('socket_error' in events,
+                "Stream error event not raised: %s" % events)
+
+
 suite = unittest.TestLoader().loadTestsFromTestCase(TestStreamTester)