diff --git a/src/Monolog/Handler/SocketHandler.php b/src/Monolog/Handler/SocketHandler.php
index c778da2f..32a9295d 100644
--- a/src/Monolog/Handler/SocketHandler.php
+++ b/src/Monolog/Handler/SocketHandler.php
@@ -5,24 +5,24 @@
 
 namespace Monolog\Handler;
 
-use Monolog\Handler\SocketHandler\Socket;
-use Monolog\Handler\SocketHandler\PersistentSocket;
 use Monolog\Logger;
+use Monolog\Handler\SocketHandler\Exception\ConnectionException;
+use Monolog\Handler\SocketHandler\Exception\WriteToSocketException;
+
 
 /**
  * Stores to any socket - uses fsockopen() or pfsockopen().
  * 
- * @see Monolog\Handler\SocketHandler\Socket
- * @see Monolog\Handler\SocketHandler\PersistentSocket
  * @see http://php.net/manual/en/function.fsockopen.php
  */
 class SocketHandler extends AbstractProcessingHandler
 {
-    /**
-     * @var Socket
-     */
-    private $socket;
- 
+    private $connectionString;
+    private $connectionTimeout;
+    private $resource;
+    private $timeout = 0;
+    private $persistent = false;
+    
     /**
      * @param string $connectionString
      * @param integer $level The minimum logging level at which this handler will be triggered
@@ -31,17 +31,21 @@ class SocketHandler extends AbstractProcessingHandler
     public function __construct($connectionString, $level = Logger::DEBUG, $bubble = true)
     {
         parent::__construct($level, $bubble);
-        $this->socket = new Socket($connectionString);
+        $this->connectionString = $connectionString;
+        $this->connectionTimeout = (float)ini_get('default_socket_timeout');
     }
-
+   
     /**
-     * Inject socket - allows you to configure timeouts.
+     * Connect (if necessary) and write to the socket
      * 
-     * @param Socket $socket 
+     * @throws Monolog\Handler\SocketHandler\Exception\ConnectionException
+     * @throws Monolog\Handler\SocketHandler\Exception\WriteToSocketException
+     * @param string $string
      */
-    public function setSocket(Socket $socket)
+    public function write(array $record)
     {
-        $this->socket = $socket;
+        $this->connectIfNotConnected();
+        $this->writeToSocket((string) $record['formatted']);
     }
     
     /**
@@ -49,17 +53,172 @@ class SocketHandler extends AbstractProcessingHandler
      */
     public function close()
     {
-        if ($this->socket instanceof PersistentSocket) {
+        if ($this->isPersistent()) {
             return;
         }
-        $this->socket->close();
+        $this->closeSocket();
+    }
+    
+    public function closeSocket()
+    {
+        if (is_resource($this->resource)) {
+            fclose($this->resource);
+            $this->resource = null;
+        }
+    }
+    
+    public function setPersistent($boolean)
+    {
+        $this->persistent = (boolean)$boolean;
+    }
+    
+    /**
+     * Set connection timeout.  Only has effect before we connect.
+     * 
+     * @see http://php.net/manual/en/function.fsockopen.php
+     * @param integer $seconds 
+     */
+    public function setConnectionTimeout($seconds)
+    {
+        $this->validateTimeout($seconds);
+        $this->connectionTimeout = (float)$seconds;
+    }
+    
+    /**
+     * Set write timeout. Only has effect before we connect.
+     * 
+     * @see http://php.net/manual/en/function.stream-set-timeout.php
+     * @param type $seconds 
+     */
+    public function setTimeout($seconds)
+    {
+        $this->validateTimeout($seconds);
+        $this->timeout = (int)$seconds;
+    }
+    
+    private function validateTimeout($value)
+    {
+        $ok = filter_var($value, FILTER_VALIDATE_INT, array('options' => array(
+            'min_range' => 0,
+        )));
+        if ($ok === false) {
+            throw new \InvalidArgumentException("Timeout must be 0 or a positive integer (got $value)");
+        }
+    }
+    
+    public function getConnectionString()
+    {
+        return $this->connectionString;
+    }
+    
+    public function isPersistent()
+    {
+        return $this->persistent;
+    }
+    
+    public function getConnectionTimeout() {
+        return $this->connectionTimeout;
+    }
+    
+    public function getTimeout() {
+        return $this->timeout;
+    }
+    
+    /**
+     * Allow injecting a resource opened somewhere else. Used in tests.
+     *
+     * @throws \InvalidArgumentException
+     * @param resource $resource 
+     */
+    public function setResource($resource)
+    {
+        if (is_resource($resource)) {
+            $this->resource = $resource;
+        } else {
+            throw new \InvalidArgumentException("Expected a resource");
+        }
     }
 
-    /**
-     * {@inheritdoc}
-     */
-    protected function write(array $record)
+    private function connectIfNotConnected()
     {
-        $this->socket->write((string) $record['formatted']);
+        if ($this->isConnected()) {
+            return;
+        }
+        $this->connect();
+    }
+    
+    /**
+     * Check to see if the socket is currently available.
+     * 
+     * UDP might appear to be connected but might fail when writing.  See http://php.net/fsockopen for details.
+     * 
+     * @return boolean
+     */
+    public function isConnected()
+    {
+        return is_resource($this->resource)
+               && !feof($this->resource);  // on TCP - other party can close connection.
+    }
+    
+    private function connect()
+    {
+        $this->createSocketResource();
+        $this->setSocketTimeout();
+    }
+    
+    protected function createSocketResource()
+    {
+        if ($this->persistent) {
+            @$resource = pfsockopen($this->connectionString, -1, $errno, $errstr, $this->connectionTimeout);
+        } else {
+            @$resource = fsockopen($this->connectionString, -1, $errno, $errstr, $this->connectionTimeout);
+        }
+        if (!$resource) {
+            throw new ConnectionException("Failed connecting to $this->connectionString ($errno: $errstr)");
+        }
+        $this->resource = $resource;
+    }
+    
+    private function setSocketTimeout()
+    {
+        if (!stream_set_timeout($this->resource, $this->timeout)) {
+            throw new ConnectionException("Failed setting timeout with stream_set_timeout()");
+        }
+    }
+    
+    protected function writeToSocket($data)
+    {
+        $length = strlen($data);
+        $sent = 0;
+        while ($this->isConnected() && $sent < $length) {
+            $chunk = $this->fwrite(substr($data, $sent));
+            if ($chunk === false) {
+                throw new WriteToSocketException("Could not write to socket");
+            }
+            $sent += $chunk;
+            $socketInfo = $this->stream_get_meta_data();
+            if ($socketInfo['timed_out']) {
+                throw new WriteToSocketException("Write timed-out");
+            }
+        }
+        if (!$this->isConnected() && $sent < $length) {
+            throw new WriteToSocketException("End-of-file reached, probably we got disconnected (sent $sent of $length)");
+        }
+    }
+    
+    /**
+     * Allow mock
+     */
+    protected function fwrite($data)
+    {
+        return @fwrite($this->resource, $data);
+    }
+    
+    /**
+     * Allow mock
+     */
+    protected function stream_get_meta_data()
+    {
+        return stream_get_meta_data($this->resource);
     }
 }
diff --git a/src/Monolog/Handler/SocketHandler/MockSocket.php b/src/Monolog/Handler/SocketHandler/MockSocket.php
deleted file mode 100644
index 2826433d..00000000
--- a/src/Monolog/Handler/SocketHandler/MockSocket.php
+++ /dev/null
@@ -1,65 +0,0 @@
-<?php
-
-/*
- * This file is part of the Monolog package.
- *
- * (c) Jordi Boggiano <j.boggiano@seld.be>
- *
- * For the full copyright and license information, please view the LICENSE
- * file that was distributed with this source code.
- */
-
-namespace Monolog\Handler\SocketHandler;
-
-use Monolog\Handler\SocketHandler\Exception\ConnectionException;
-use Monolog\Handler\SocketHandler\Exception\WriteToSocketException;
-
-class MockSocket extends Socket
-{
-    private $connectTimeoutMock = 0;
-    private $timeoutMock = 0;
-    
-    
-    public function __construct($connectionString)
-    {
-        if (is_resource($connectionString)) {
-            $this->resource = $connectionString;
-        } else {
-            $this->connectionString = $connectionString;
-        }
-    }
-    
-    public function setFailConnectionTimeout($seconds)
-    {
-        $this->connectTimeoutMock = (int)$seconds;
-    }
-    
-    public function setFailTimeout($seconds)
-    {
-        $this->timeoutMock = (int)$seconds;
-    }
-    
-    protected function createSocketResource()
-    {
-        if ($this->connectTimeoutMock > 0) {
-            throw new ConnectionException("Mocked connection timeout");
-        }
-        $this->resource = fopen('php://memory', '+a');
-    }
-    
-    protected function writeToSocket($data) {
-        if ($this->timeoutMock > 0) {
-            throw new WriteToSocketException("Mocked write timeout");
-        }
-        return parent::writeToSocket($data);
-    }
-    
-    protected function setSocketTimeout()
-    {
-        // php://memory does not support this
-    }
-    
-    public function getResource() {
-        return $this->resource;
-    }
-}
diff --git a/src/Monolog/Handler/SocketHandler/PersistentSocket.php b/src/Monolog/Handler/SocketHandler/PersistentSocket.php
deleted file mode 100644
index c4a2d788..00000000
--- a/src/Monolog/Handler/SocketHandler/PersistentSocket.php
+++ /dev/null
@@ -1,32 +0,0 @@
-<?php
-
-/*
- * This file is part of the Monolog package.
- *
- * (c) Jordi Boggiano <j.boggiano@seld.be>
- *
- * For the full copyright and license information, please view the LICENSE
- * file that was distributed with this source code.
- */
-
-namespace Monolog\Handler\SocketHandler;
-
-use Monolog\Handler\SocketHandler\Exception\ConnectionException;
-
-/**
- * Same as Socket but uses pfsockopen() instead allowing the connection to be reused in other requests.
- * 
- * @see http://php.net/manual/en/function.pfsockopen.php
- * @author Pablo de Leon Belloc <pablolb@gmail.com>
- */
-class PersistentSocket extends Socket
-{
-    protected function createSocketResource()
-    {
-        @$resource = pfsockopen($this->connectionString, -1, $errno, $errstr, $this->connectionTimeout);
-        if (!$resource) {
-            throw new ConnectionException("Failed connecting to $this->connectionString ($errno: $errstr)");
-        }
-        $this->resource = $resource;
-    }
-}
diff --git a/src/Monolog/Handler/SocketHandler/Socket.php b/src/Monolog/Handler/SocketHandler/Socket.php
deleted file mode 100644
index 281d162f..00000000
--- a/src/Monolog/Handler/SocketHandler/Socket.php
+++ /dev/null
@@ -1,185 +0,0 @@
-<?php
-
-/*
- * This file is part of the Monolog package.
- *
- * (c) Jordi Boggiano <j.boggiano@seld.be>
- *
- * For the full copyright and license information, please view the LICENSE
- * file that was distributed with this source code.
- */
-
-namespace Monolog\Handler\SocketHandler;
-
-use Monolog\Handler\SocketHandler\Exception\ConnectionException;
-use Monolog\Handler\SocketHandler\Exception\WriteToSocketException;
-
-/**
- * Small class which writes to a socket.
- * Timeout settings must be set before first write to have any effect.
- * 
- * @see http://php.net/manual/en/function.fsockopen.php
- * @author Pablo de Leon Belloc <pablolb@gmail.com>
- */
-class Socket
-{
-    protected $connectionString;
-    protected $connectionTimeout;
-    protected $resource;
-    private $timeout = 0;
-    
-    /**
-     * @param string $connectionString As interpreted by fsockopen()
-     */
-    public function __construct($connectionString)
-    {
-        $this->connectionString = $connectionString;
-        $this->connectionTimeout = (float)ini_get('default_socket_timeout');
-    }
-    
-    public function getConnectionString()
-    {
-        return $this->connectionString;
-    }
-    
-    /**
-     * Set connection timeout.  Only has effect before we connect.
-     * 
-     * @see http://php.net/manual/en/function.fsockopen.php
-     * @param integer $seconds 
-     */
-    public function setConnectionTimeout($seconds)
-    {
-        $this->validateTimeout($seconds);
-        $this->connectionTimeout = (float)$seconds;
-    }
-    
-    /**
-     * Set write timeout. Only has effect before we connect.
-     * 
-     * @see http://php.net/manual/en/function.stream-set-timeout.php
-     * @param type $seconds 
-     */
-    public function setTimeout($seconds)
-    {
-        $this->validateTimeout($seconds);
-        $this->timeout = (int)$seconds;
-    }
-    
-    private function validateTimeout($value)
-    {
-        $ok = filter_var($value, FILTER_VALIDATE_INT, array('options' => array(
-            'min_range' => 0,
-        )));
-        if ($ok === false) {
-            throw new \InvalidArgumentException("Timeout must be 0 or a positive integer (got $value)");
-        }
-    }
-    
-    public function getConnectionTimeout() {
-        return $this->connectionTimeout;
-    }
-    
-    public function getTimeout() {
-        return $this->timeout;
-    }
-    
-    public function close()
-    {
-        if (is_resource($this->resource)) {
-            fclose($this->resource);
-            $this->resource = null;
-        }
-    }
-    
-    /**
-     * Allow injecting a resource opened somewhere else. Used in tests.
-     *
-     * @throws \InvalidArgumentException
-     * @param resource $resource 
-     */
-    public function setResource($resource)
-    {
-        if (is_resource($resource)) {
-            $this->resource = $resource;
-        } else {
-            throw new \InvalidArgumentException("Expected a resource");
-        }
-    }
-
-    /**
-     * Connect (if necessary) and write to the socket
-     * 
-     * @throws Monolog\Handler\SocketHandler\Exception\ConnectionException
-     * @throws Monolog\Handler\SocketHandler\Exception\WriteToSocketException
-     * @param string $string
-     */
-    public function write($string)
-    {
-        $this->connectIfNotConnected();
-        $this->writeToSocket($string);
-    }
-    
-    protected function connectIfNotConnected()
-    {
-        if ($this->isConnected()) {
-            return;
-        }
-        $this->connect();
-    }
-    
-    /**
-     * Check to see if the socket is currently available.
-     * 
-     * UDP might appear to be connected but might fail when writing.  See http://php.net/fsockopen for details.
-     * 
-     * @return boolean
-     */
-    public function isConnected()
-    {
-        return is_resource($this->resource)
-               && !feof($this->resource);  // on TCP - other party can close connection.
-    }
-    
-    protected function connect()
-    {
-        $this->createSocketResource();
-        $this->setSocketTimeout();
-    }
-    
-    protected function createSocketResource()
-    {
-        @$resource = fsockopen($this->connectionString, -1, $errno, $errstr, $this->connectionTimeout);
-        if (!$resource) {
-            throw new ConnectionException("Failed connecting to $this->connectionString ($errno: $errstr)");
-        }
-        $this->resource = $resource;
-    }
-    
-    protected function setSocketTimeout()
-    {
-        if (!stream_set_timeout($this->resource, $this->timeout)) {
-            throw new ConnectionException("Failed setting timeout with stream_set_timeout()");
-        }
-    }
-    
-    protected function writeToSocket($data)
-    {
-        $length = strlen($data);
-        $sent = 0;
-        while ($this->isConnected() && $sent < $length) {
-            @$chunk = fwrite($this->resource, substr($data, $sent));
-            if ($chunk === false) {
-                throw new WriteToSocketException("Could not write to socket");
-            }
-            $sent += $chunk;
-            $socketInfo = stream_get_meta_data($this->resource);
-            if ($socketInfo['timed_out']) {
-                throw new WriteToSocketException("Write timed-out");
-            }
-        }
-        if (!$this->isConnected() && $sent < $length) {
-            throw new WriteToSocketException("End-of-file reached, probably we got disconnected (sent $sent of $length)");
-        }
-    }
-}
diff --git a/tests/Monolog/Handler/SocketHandler/SocketTest.php b/tests/Monolog/Handler/SocketHandler/SocketTest.php
index cc5b237a..a7b1d49f 100644
--- a/tests/Monolog/Handler/SocketHandler/SocketTest.php
+++ b/tests/Monolog/Handler/SocketHandler/SocketTest.php
@@ -1,124 +1,8 @@
 <?php
 
-/*
- * This file is part of the Monolog package.
- *
- * (c) Jordi Boggiano <j.boggiano@seld.be>
- *
- * For the full copyright and license information, please view the LICENSE
- * file that was distributed with this source code.
- */
-
-
 namespace Monolog\Handler\SocketHandler;
 
 class SocketTest extends \PHPUnit_Framework_TestCase
 {
     
-    /**
-     * @expectedException Monolog\Handler\SocketHandler\Exception\ConnectionException
-     */
-    public function testInvalidHostname() {
-        $socket = new Socket('garbage://here');
-        $socket->write('data');
-    }
-    
-    /**
-     * @expectedException \InvalidArgumentException
-     */
-    public function testBadConnectionTimeout()
-    {
-        $socket = new Socket('localhost:1234');
-        $socket->setConnectionTimeout(-1);
-    }
-    
-    public function testSetConnectionTimeout()
-    {
-        $socket = new Socket('localhost:1234');
-        $socket->setConnectionTimeout(10);
-        $this->assertEquals(10, $socket->getConnectionTimeout());
-    }
-    
-    /**
-     * @expectedException \InvalidArgumentException
-     */
-    public function testBadTimeout()
-    {
-        $socket = new Socket('localhost:1234');
-        $socket->setTimeout(-1);
-    }
-    
-    public function testSetTimeout()
-    {
-        $socket = new Socket('localhost:1234');
-        $socket->setTimeout(10);
-        $this->assertEquals(10, $socket->getTimeout());
-    }
-    
-    public function testSetConnectionString()
-    {
-        $socket = new Socket('tcp://localhost:9090');
-        $this->assertEquals('tcp://localhost:9090', $socket->getConnectionString());
-    }
-    
-    public function testConnectionRefuesed()
-    {
-        try {
-            $socket = new Socket('127.0.0.1:7894');
-            $socket->setTimeout(1);
-            $string = 'Hello world';
-            $socket->write($string);
-            $this->fail("Shoul not connect - are you running a server on 127.0.0.1:7894 ?");
-        } catch (\Monolog\Handler\SocketHandler\Exception\ConnectionException $e) {
-        }
-    }
-    
-    /**
-     * @expectedException Monolog\Handler\SocketHandler\Exception\ConnectionException
-     */
-    public function testConnectionTimeoutWithMock()
-    {
-        $socket = new MockSocket('localhost:54321');
-        $socket->setConnectionTimeout(10);
-        $socket->setFailConnectionTimeout(10);
-        $socket->write('Hello world');
-    }
-    
-    /**
-     * @expectedException Monolog\Handler\SocketHandler\Exception\WriteToSocketException
-     */
-    public function testWriteTimeoutWithMock()
-    {
-        $socket = new MockSocket('localhost:54321');
-        $socket->setTimeout(10);
-        $socket->setFailTimeout(10);
-        $socket->write('Hello world');
-    }
-    
-    public function testWriteWithMock()
-    {
-        $socket = new MockSocket('localhost:54321');
-        $socket->write('Hello world');
-        $res = $socket->getResource();
-        fseek($res, 0);
-        $this->assertEquals('Hello world', fread($res, 1024));
-    }
-    
-    public function testClose()
-    {
-        $resource = fopen('php://memory', 'a+');
-        $socket = new MockSocket($resource);
-        $this->assertTrue(is_resource($resource));
-        $socket->close();
-        $this->assertFalse(is_resource($resource));
-    }
-
-    /**
-     * @expectedException \InvalidArgumentException
-     */
-    public function testInjectBadResourceThrowsException()
-    {
-        $socket = new Socket('');
-        $socket->setResource('');
-    }
 }
\ No newline at end of file
diff --git a/tests/Monolog/Handler/SocketHandlerTest.php b/tests/Monolog/Handler/SocketHandlerTest.php
index cc5c9201..97a69e09 100644
--- a/tests/Monolog/Handler/SocketHandlerTest.php
+++ b/tests/Monolog/Handler/SocketHandlerTest.php
@@ -11,54 +11,250 @@
 
 namespace Monolog\Handler;
 
-use Monolog\Handler\SocketHandler\MockSocket;
-use Monolog\Handler\SocketHandler\Socket;
-use Monolog\Handler\SocketHandler\PersistentSocket;
-
 use Monolog\TestCase;
 use Monolog\Logger;
-
-
+use Monolog\Handler\SocketHandler\Exception\ConnectionException;
+use Monolog\Handler\SocketHandler\Exception\WriteToSocketException;
 
 /**
  * @author Pablo de Leon Belloc <pablolb@gmail.com>
  */
 class SocketHandlerTest extends TestCase
 {
-    public function testWrite()
-    {
-        $socket = new MockSocket('localhost');
-        $handler = new SocketHandler('localhost');
-        $handler->setSocket($socket);
-        $handler->setFormatter($this->getIdentityFormatter());
-        $handler->handle($this->getRecord(Logger::WARNING, 'test'));
-        $handler->handle($this->getRecord(Logger::WARNING, 'test2'));
-        $handler->handle($this->getRecord(Logger::WARNING, 'test3'));
-        $handle = $socket->getResource();
-        fseek($handle, 0);
-        $this->assertEquals('testtest2test3', fread($handle, 100));
+    /**
+     * @var Monolog\Handler\SocketHandler
+     */
+    private $handler;
+    
+    /**
+     * @var resource
+     */
+    private $res;
+    
+    /**
+     * @expectedException Monolog\Handler\SocketHandler\Exception\ConnectionException
+     */
+    public function testInvalidHostname() {
+        $this->createHandler('garbage://here');
+        $this->writeRecord('data');
     }
     
-    public function testCloseClosesNonPersistentSocket()
+    /**
+     * @expectedException \InvalidArgumentException
+     */
+    public function testBadConnectionTimeout()
     {
-        $socket = new Socket('localhost');
-        $res = fopen('php://memory', 'a');
-        $socket->setResource($res);
-        $handler = new SocketHandler('localhost');
-        $handler->setSocket($socket);
-        $handler->close();
-        $this->assertFalse($socket->isConnected());
+        $this->createHandler('localhost:1234');
+        $this->handler->setConnectionTimeout(-1);
+    }
+    
+    public function testSetConnectionTimeout()
+    {
+        $this->createHandler('localhost:1234');
+        $this->handler->setConnectionTimeout(10);
+        $this->assertEquals(10, $this->handler->getConnectionTimeout());
+    }
+    
+    /**
+     * @expectedException \InvalidArgumentException
+     */
+    public function testBadTimeout()
+    {
+        $this->createHandler('localhost:1234');
+        $this->handler->setTimeout(-1);
+    }
+    
+    public function testSetTimeout()
+    {
+        $this->createHandler('localhost:1234');
+        $this->handler->setTimeout(10);
+        $this->assertEquals(10, $this->handler->getTimeout());
+    }
+    
+    public function testSetConnectionString()
+    {
+        $this->createHandler('tcp://localhost:9090');
+        $this->assertEquals('tcp://localhost:9090', $this->handler->getConnectionString());
+    }
+    
+    public function testConnectionRefuesed()
+    {
+        try {
+            $this->createHandler('127.0.0.1:7894');
+            $string = 'Hello world';
+            $this->writeRecord($string);
+            $this->fail("Shoul not connect - are you running a server on 127.0.0.1:7894 ?");
+        } catch (\Monolog\Handler\SocketHandler\Exception\ConnectionException $e) {
+        }
+    }
+    
+    /**
+     * @expectedException Monolog\Handler\SocketHandler\Exception\ConnectionException
+     */
+    public function testConnectionTimeoutWithMock()
+    {
+        $this->setMockHandler(array('createSocketResource'));
+        $this->handler->expects($this->once())
+                      ->method('createSocketResource')
+                      ->will($this->throwException(new ConnectionException()));
+        $this->writeRecord('Hello world');
+    }
+    
+    /**
+     * @expectedException Monolog\Handler\SocketHandler\Exception\WriteToSocketException
+     */
+    public function testWriteFailsOnIfFwriteReturnsFalse()
+    {
+        $this->setMockHandler(array('fwrite'));
+        
+        $map = array(
+            array('Hello world', 6),
+            array('world', false),
+        );
+        
+        $this->handler->expects($this->exactly(2))
+                      ->method('fwrite')
+                      ->will($this->returnValueMap($map));
+        
+        $this->injectMemoryResource();
+        $this->writeRecord('Hello world');
+    }
+    
+    /**
+     * @expectedException Monolog\Handler\SocketHandler\Exception\WriteToSocketException
+     */
+    public function testWriteFailsIfStreamTimesOut()
+    {
+        $this->setMockHandler(array('fwrite', 'stream_get_meta_data'));
+        
+        $map = array(
+            array('Hello world', 6),
+            array('world', 5),
+        );
+        
+        $this->handler->expects($this->exactly(1))
+                      ->method('fwrite')
+                      ->will($this->returnValueMap($map));
+        $this->handler->expects($this->exactly(1))
+                      ->method('stream_get_meta_data')
+                      ->will($this->returnValue(array('timed_out' => true)));
+        
+        
+        $this->injectMemoryResource();
+        $this->writeRecord('Hello world');
+    }
+    
+    /**
+     * @expectedException Monolog\Handler\SocketHandler\Exception\WriteToSocketException
+     */
+    public function testWriteFailsOnIncompleteWrite()
+    {
+        $this->setMockHandler(array('fwrite', 'isConnected'));
+        
+        $map = array(
+            array('Hello world', 6),
+            array('world', 5),
+        );
+        
+        $this->handler->expects($this->exactly(1))
+                      ->method('fwrite')
+                      ->will($this->returnValueMap($map));
+        $this->handler->expects($this->at(0))
+                      ->method('isConnected')
+                      ->will($this->returnValue(true));
+        $this->handler->expects($this->at(1))
+                      ->method('isConnected')
+                      ->will($this->returnValue(true));
+        $this->handler->expects($this->at(2))
+                      ->method('isConnected')
+                      ->will($this->returnValue(false));
+        
+        $this->injectMemoryResource();
+        $this->writeRecord('Hello world');
+    }
+    
+    public function testWriteWithMemoryFile()
+    {
+        $this->createHandler('localhost:54321');
+        $this->injectMemoryResource();
+        $this->writeRecord('test1');
+        $this->writeRecord('test2');
+        $this->writeRecord('test3');
+        fseek($this->res, 0);
+        $this->assertEquals('test1test2test3', fread($this->res, 1024));
+    }
+    
+    public function testWriteWithMock()
+    {
+        $this->setMockHandler(array('fwrite'));
+        
+        $map = array(
+            array('Hello world', 6),
+            array('world', 5),
+        );
+        
+        $this->handler->expects($this->exactly(2))
+                      ->method('fwrite')
+                      ->will($this->returnValueMap($map));
+        
+        $this->injectMemoryResource();
+        $this->writeRecord('Hello world');
+    }
+    
+    public function testClose()
+    {
+        $this->createHandler('localhost:54321');
+        $this->injectMemoryResource();
+        $this->writeRecord('Hello world');
+        $this->assertTrue(is_resource($this->res));
+        $this->handler->close();
+        $this->assertFalse(is_resource($this->res));
     }
     
     public function testCloseDoesNotClosePersistentSocket()
     {
-        $socket = new PersistentSocket('localhost');
-        $res = fopen('php://memory', 'a');
-        $socket->setResource($res);
-        $handler = new SocketHandler('localhost');
-        $handler->setSocket($socket);
-        $handler->close();
-        $this->assertTrue($socket->isConnected());
+        $this->createHandler('localhost:54321');
+        $this->handler->setPersistent(true);
+        $this->injectMemoryResource();
+        $this->writeRecord('Hello world');
+        $this->assertTrue(is_resource($this->res));
+        $this->handler->close();
+        $this->assertTrue(is_resource($this->res));
+    }
+
+    /**
+     * @expectedException \InvalidArgumentException
+     */
+    public function testInjectBadResourceThrowsException()
+    {
+        $this->createHandler('');
+        $this->handler->setResource('');
     }
     
+    private function createHandler($connectionString)
+    {
+        $this->handler = new SocketHandler($connectionString);
+        $this->handler->setFormatter($this->getIdentityFormatter());
+    }
+    
+    private function writeRecord($string)
+    {
+        $this->handler->handle($this->getRecord(Logger::WARNING, $string));
+    }
+    
+    private function injectMemoryResource()
+    {
+        $this->res = fopen('php://memory', 'a');
+        $this->handler->setResource($this->res);
+    }
+    
+    private function setMockHandler(array $methods)
+    {
+        $this->handler = $this->getMock(
+                '\Monolog\Handler\SocketHandler',
+                $methods,
+                array('localhost:1234')
+        );
+        $this->handler->setFormatter($this->getIdentityFormatter());
+    }
 }