Introduce the mock module to improve unit tests.

There are a few socket corner-cases that are incredibly hard to reproduce in a
unit test suite, so we use mock to trigger the exceptions instead.

1 files changed, 18 insertions, 1 deletions

diff --git a/test/test_tcp.py b/test/test_tcp.py
index ad09143d..e7524fdc 100644
--- a/test/test_tcp.py
+++ b/test/test_tcp.py
@@ -1,5 +1,6 @@
-import cStringIO, threading, Queue, time
+import cStringIO, threading, Queue, time, socket
 from netlib import tcp, certutils, test
+import mock
 import tutils
 
 class SNIHandler(tcp.BaseHandler):
@@ -275,6 +276,22 @@ class TestFileLike:
         s.write("x")
         assert s.get_log() == "xx"
 
+    def test_writer_flush_error(self):
+        s = cStringIO.StringIO()
+        s = tcp.Writer(s)
+        o = mock.MagicMock()
+        o.flush = mock.MagicMock(side_effect=socket.error)
+        s.o = o
+        tutils.raises(tcp.NetLibDisconnect, s.flush)
+
+    def test_reader_read_error(self):
+        s = cStringIO.StringIO("foobar\nfoobar")
+        s = tcp.Reader(s)
+        o = mock.MagicMock()
+        o.read = mock.MagicMock(side_effect=socket.error)
+        s.o = o
+        tutils.raises(tcp.NetLibDisconnect, s.read, 10)
+
     def test_reset_timestamps(self):
         s = cStringIO.StringIO("foobar\nfoobar")
         s = tcp.Reader(s)