diff --git a/tests/Woketo/Message/SimpleMessageHandlerTest.php b/tests/Woketo/Message/SimpleMessageHandlerTest.php
index 221540d..f83f3dd 100644
--- a/tests/Woketo/Message/SimpleMessageHandlerTest.php
+++ b/tests/Woketo/Message/SimpleMessageHandlerTest.php
@@ -43,6 +43,16 @@ public function testItEchosOnError()
 
         $this->assertContains('foobar', $out);
     }
+
+    public function testItDoNothingOnDisconnection()
+    {
+        \ob_start();
+        $res = $this->instance->onDisconnect($this->prophesize(Connection::class)->reveal());
+        $out = \ob_get_clean();
+
+        $this->assertNull($res);
+        $this->assertEquals('', $out);
+    }
 }
 
 class SimpleMessageHandlerImplementation extends SimpleMessageHandler
diff --git a/tests/Woketo/Server/ConnectionTest.php b/tests/Woketo/Server/ConnectionTest.php
index 86ac0f3..062cd64 100644
--- a/tests/Woketo/Server/ConnectionTest.php
+++ b/tests/Woketo/Server/ConnectionTest.php
@@ -11,6 +11,7 @@
 
 namespace Test\Woketo\Server;
 
+use Evenement\EventEmitterTrait;
 use Nekland\Woketo\Message\MessageHandlerInterface;
 use Nekland\Woketo\Rfc6455\Frame;
 use Nekland\Woketo\Rfc6455\Handshake\ServerHandshake;
@@ -88,6 +89,26 @@ public function testItSupportsBinaryMessage()
         $reactMock->emit('data', [$binaryFrame]);
     }
 
+    public function testItCallOnDisconnectOnHandlerWhenDisconnect()
+    {
+        // Mocks
+        $reactMock = new ReactConnectionMock();
+        $handler = $this->prophesize(MessageHandlerInterface::class);
+        $loop = $this->prophesize(LoopInterface::class);
+        /** @var MessageProcessor $processor */
+        $processor = $this->prophesize(MessageProcessor::class);
+        $handshakeProcessor = $this->prophesize(ServerHandshake::class);
+
+        // Init
+        $connection = new Connection($reactMock, function () use ($handler) {return $handler->reveal();}, $loop->reveal(), $processor->reveal(), $handshakeProcessor->reveal());
+        $server = new ReactConnectionMock();
+        $server->emit('connection', [$connection]);
+
+
+        $handler->onDisconnect(Argument::type(Connection::class))->shouldBeCalled();
+        $reactMock->emit('end');
+    }
+
     private function getHandshake()
     {
         return "GET /foo HTTP/1.1\r\n"
@@ -109,32 +130,10 @@ private function getHandshake()
 
 class ReactConnectionMock implements ConnectionInterface
 {
-    public function __construct()
-    {
-    }
-
-    private $on = [];
-
-    public function on($event, callable $listener)
-    {
-        $this->on[$event] = $listener;
-    }
-
-    public function emit($event, array $arguments = [])
-    {
-        call_user_func_array($this->on[$event], $arguments);
-    }
+    use EventEmitterTrait;
 
     public function getRemoteAddress() {}
 
-    public function once($event, callable $listener) {}
-
-    public function removeListener($event, callable $listener) {}
-
-    public function removeAllListeners($event = null) {}
-
-    public function listeners($event = null) {}
-
     public function isReadable(){}
 
     public function pause() {}
diff --git a/tests/Woketo/Server/WebSocketServerTest.php b/tests/Woketo/Server/WebSocketServerTest.php
index fbc44f3..690bbfa 100644
--- a/tests/Woketo/Server/WebSocketServerTest.php
+++ b/tests/Woketo/Server/WebSocketServerTest.php
@@ -11,6 +11,7 @@
 namespace Test\Woketo\Server;
 
 
+use Evenement\EventEmitterTrait;
 use Nekland\Woketo\Core\AbstractConnection;
 use Nekland\Woketo\Exception\ConfigException;
 use Nekland\Woketo\Exception\RuntimeException;
@@ -83,6 +84,31 @@ public function onConnection(AbstractConnection $connection)
         $this->assertTrue($handler->called);
     }
 
+    public function testItCallTheDisconnectionMethodOfHandler()
+    {
+        $handler = new class extends TextMessageHandler {
+            public $called = false;
+            public function onMessage(string $data, AbstractConnection $connection) {}
+            public function onConnection(AbstractConnection $connection) {}
+
+            public function onDisconnect(AbstractConnection $connection)
+            {
+                $this->called = true;
+            }
+        };
+
+        $server = new WebSocketServer(1000, '127.0.0.1', ['prod' => false]);
+        $server->setMessageHandler($handler);
+        $server->setLoop($this->prophesize(LoopInterface::class)->reveal());
+        $server->setSocketServer($socket = new FakeSocketServerForTestMethodHandlerConnection());
+        $server->setLogger(new NullLogger());
+        $server->start();
+        $socket->callCb($co = new ServerReactConnectionMock());
+        $co->emit('data', [self::getHandshake()]);
+        $co->emit('end');
+        $this->assertTrue($handler->called);
+    }
+
     /**
      * @dataProvider getMessageHandlerTestData
      */
@@ -275,32 +301,10 @@ public function close() {}
 
 class ServerReactConnectionMock implements ConnectionInterface
 {
-    public function __construct()
-    {
-    }
-
-    private $on = [];
-
-    public function on($event, callable $listener)
-    {
-        $this->on[$event] = $listener;
-    }
-
-    public function emit($event, array $arguments = [])
-    {
-        call_user_func_array($this->on[$event], $arguments);
-    }
+    use EventEmitterTrait;
 
     public function getRemoteAddress() {}
 
-    public function once($event, callable $listener) {}
-
-    public function removeListener($event, callable $listener) {}
-
-    public function removeAllListeners($event = null) {}
-
-    public function listeners($event = null) {}
-
     public function isReadable(){}
 
     public function pause() {}