diff --git a/QtPyNetwork/__init__.py b/QtPyNetwork/__init__.py
index 43c4ab0..49e0fc1 100644
--- a/QtPyNetwork/__init__.py
+++ b/QtPyNetwork/__init__.py
@@ -1 +1 @@
-__version__ = "0.6.1"
+__version__ = "0.7.0"
diff --git a/QtPyNetwork2/balancers/AbstractBalancer.py b/QtPyNetwork/balancer/AbstractBalancer.py
similarity index 100%
rename from QtPyNetwork2/balancers/AbstractBalancer.py
rename to QtPyNetwork/balancer/AbstractBalancer.py
diff --git a/QtPyNetwork2/balancers/NoBalancer.py b/QtPyNetwork/balancer/NoBalancer.py
similarity index 96%
rename from QtPyNetwork2/balancers/NoBalancer.py
rename to QtPyNetwork/balancer/NoBalancer.py
index d20e978..d557d64 100644
--- a/QtPyNetwork2/balancers/NoBalancer.py
+++ b/QtPyNetwork/balancer/NoBalancer.py
@@ -105,10 +105,7 @@ def write(self, client_id: int, data: bytes):
                 data = pack(HEADER, len(data)) + data
                 socket.write(data)
                 socket.flush()
-                print(f"Written {len(data)} bytes to {socket.objectName()}")
                 return
-        print(f"{client_id} - {data}")
-        print([socket.objectName() for socket in self.sockets])
         self.client_error.emit(client_id, Exception(f"Client {client_id} not found"))
 
     @Slot(bytes)
diff --git a/QtPyNetwork2/balancers/ThreadBalancer.py b/QtPyNetwork/balancer/ThreadBalancer.py
similarity index 50%
rename from QtPyNetwork2/balancers/ThreadBalancer.py
rename to QtPyNetwork/balancer/ThreadBalancer.py
index 4fb0b53..8bff8d4 100644
--- a/QtPyNetwork2/balancers/ThreadBalancer.py
+++ b/QtPyNetwork/balancer/ThreadBalancer.py
@@ -2,21 +2,21 @@
 from qtpy.QtNetwork import QAbstractSocket, QUdpSocket
 
 import logging
-from struct import unpack, calcsize
 
-from .AbstractBalancer import AbstractBalancer, HEADER, HEADER_SIZE
+from QtPyNetwork.common import read, write
+from .AbstractBalancer import AbstractBalancer
 
 
 class _Worker(QObject):
-    disconnected = Signal()
-    connected = Signal(str, int)
-    readyRead = Signal(bytes)
-    error = Signal(Exception)
-    closed = Signal()
+    disconnected = Signal(int)
+    connected = Signal(int, str, int)
+    ready_read = Signal(int, bytes)
+    error = Signal(int, Exception)
+    closed = Signal(int)
 
     close_signal = Signal()
 
-    write = Signal(bytes)
+    write_signal = Signal(bytes)
 
     def __init__(self, client_id, socket_type: type, socket_descriptor: int):
         super(_Worker, self).__init__()
@@ -29,6 +29,9 @@ def __init__(self, client_id, socket_type: type, socket_descriptor: int):
         self.size_left = 0
         self.data = b""
 
+        self.close_signal.connect(self.__on_close_signal)
+        self.write_signal.connect(self.__on_write_signal)
+
     @Slot()
     def start(self):
         socket: QAbstractSocket = self.socket_type()
@@ -50,49 +53,35 @@ def __on_socket_ready_read(self):
         Note:
             Emits message signal.
         """
-
-        while self.socket.bytesAvailable():
-            if self.size_left > 0:
-                data = self.socket.read(self.size_left)
-                self.size_left = self.size_left - len(data)
-                if self.size_left > 0:
-                    self.data = self.data + data
-                else:
-                    self.readyRead.emit(self.client_id, self.data)
-                    self.data = b""
-                    self.size_left = 0
-            else:
-                header = self.socket.read(HEADER_SIZE)
-                data_size = unpack(HEADER, header)[0]
-                data = self.socket.read(data_size)
-                if len(data) < data_size:
-                    self.data = data
-                    self.size_left = data_size - len(data)
-                else:
-                    self.readyRead.emit(self.client_id, data)
-            
-            # if self.client_id in self.data:
-            #     size_left = self.data.get(client_id).get("size_left")
-            #     data = socket.read(size_left)
-            #     size_left = size_left - len(data)
-            #     if size_left > 0:
-            #         self.data[client_id]["size_left"] = size_left
-            #         self.data[client_id]["data"] += data
-            #     else:
-            #         data = self.data.get(client_id).get("data") + data
-            #         del self.data[client_id]
-            #         self.message.emit(client_id, data)
-            #
-            # else:
-            #     header = socket.read(HEADER_SIZE)
-            #     data_size = unpack(HEADER, header)[0]
-            #     message = socket.read(data_size)
-            #
-            #     if len(message) < data_size:
-            #         data_size = data_size - len(message)
-            #         self.data[client_id] = {"data": message, "size_left": data_size}
-            #     else:
-            #         self.message.emit(client_id, message)
+        data, size_left = read(self.socket, self.data, self.size_left)
+        if size_left == 0:
+            self._size_left = 0
+            self._data = b""
+            self.ready_read.emit(self.client_id, data)
+        else:
+            self._data = data
+            self._size_left = size_left
+
+
+        # while self.socket.bytesAvailable():
+        #     if self.size_left > 0:
+        #         data = self.socket.read(self.size_left)
+        #         self.size_left = self.size_left - len(data)
+        #         if self.size_left > 0:
+        #             self.data = self.data + data
+        #         else:
+        #             self.ready_read.emit(self.client_id, self.data)
+        #             self.data = b""
+        #             self.size_left = 0
+        #     else:
+        #         header = self.socket.read(HEADER_SIZE)
+        #         data_size = unpack(HEADER, header)[0]
+        #         data = self.socket.read(data_size)
+        #         if len(data) < data_size:
+        #             self.data = data
+        #             self.size_left = data_size - len(data)
+        #         else:
+        #             self.ready_read.emit(self.client_id, data)
 
     @Slot()
     def __on_socket_disconnected(self):
@@ -123,24 +112,38 @@ def __on_socket_error(self):
         error = self.socket.errorString()
         self.error.emit(self.client_id, Exception(error))
 
+    @Slot()
+    def __on_close_signal(self):
+        """Close socket.
 
-class ThreadBalancer(AbstractBalancer):
+        Note:
+            Emits closed signal.
+        """
+        try:
+            self.socket.close()
+        except RuntimeError:
+            pass
+        self.closed.emit(self.client_id)
 
     @Slot(bytes)
-    def write_all(self, message: bytes):
-        pass
+    def __on_write_signal(self, data: bytes):
+        """Write data to socket.
 
-    @Slot(int)
-    def disconnect(self, client_id: int):
-        pass
+        Args:
+            data (bytes): Data to write.
 
-    @Slot()
-    def close(self):
-        pass
+        Note:
+            Emits written signal.
+        """
+        write(self.socket, data)
+        self.socket.flush()
+
+        # data = pack(HEADER, len(data)) + data
+        # self.socket.write(data)
+        # self.socket.flush()
 
-    @Slot(int, bytes)
-    def write(self, client_id: int, message: bytes):
-        pass
+
+class ThreadBalancer(AbstractBalancer):
 
     def __init__(self):
         super(ThreadBalancer, self).__init__()
@@ -153,40 +156,48 @@ def balance(self, socket_type: type, socket_descriptor: int):
         worker = _Worker(client_id, socket_type, socket_descriptor)
         worker.setObjectName(str(client_id))
         # worker.connected.connect(self.__on_worker_socket_connected)
-        # worker.readyRead.connect(self.__on_worker_socket_readyRead)
+        # worker.ready_read.connect(self.__on_worker_socket_readyRead)
         # worker.disconnected.connect()
         worker.connected.connect(self.connected.emit)
         worker.disconnected.connect(self.disconnected.emit)
-        worker.readyRead.connect(self.message.emit)
+        worker.ready_read.connect(self.message.emit)
         worker.error.connect(self.client_error.emit)
 
         thread = QThread()
         worker.moveToThread(thread)
-        thread.started.connect(worker.run)
+        thread.started.connect(worker.start)
         self.workers.append((worker, thread))
         thread.start()
 
-    # @Slot(str, int)
-    # def __on_worker_socket_connected(self, ip: str, port: int):
-    #     client_id = int(self.sender().objectName())
-    #     self.connected.emit(client_id, ip, port)
-    #
-    # @Slot(bytes)
-    # def __on_worker_socket_readyRead(self, data: bytes):
-    #     client_id = int(self.sender().objectName())
-    #     self.connected.emit(client_id, data)
-    #
-    # @Slot()
-    # def __on_worker_socket_disconnected(self):
-    #     client_id = int(self.sender().objectName())
-    #     self.disconnected.emit(client_id)
-
-
-
+    @Slot(int, bytes)
+    def write(self, client_id: int, message: bytes):
+        worker = self.__get_worker_by_client_id(client_id)
+        print(f"Writing to {client_id}: {message}")
+        if worker:
+            worker.write_signal.emit(message)
+        else:
+            self.client_error.emit(client_id, Exception("Client not found"))
 
+    @Slot(bytes)
+    def write_all(self, message: bytes):
+        for worker in self.workers:
+            worker.write_signal.emit(message)
 
+    @Slot(int)
+    def disconnect(self, client_id: int):
+        worker = self.__get_worker_by_client_id(client_id)
+        if worker:
+            worker.close_signal.emit()
 
+    @Slot()
+    def close(self):
+        pass
 
+    @Slot(int)
+    def __get_worker_by_client_id(self, client_id: int):
+        for worker, thread in self.workers:
+            if worker.objectName() == str(client_id):
+                return worker
 
 
 
diff --git a/QtPyNetwork2/balancers/ThreadPoolBalancer.py b/QtPyNetwork/balancer/ThreadPoolBalancer.py
similarity index 100%
rename from QtPyNetwork2/balancers/ThreadPoolBalancer.py
rename to QtPyNetwork/balancer/ThreadPoolBalancer.py
diff --git a/QtPyNetwork2/balancers/__init__.py b/QtPyNetwork/balancer/__init__.py
similarity index 65%
rename from QtPyNetwork2/balancers/__init__.py
rename to QtPyNetwork/balancer/__init__.py
index f850418..5034ea7 100644
--- a/QtPyNetwork2/balancers/__init__.py
+++ b/QtPyNetwork/balancer/__init__.py
@@ -1,2 +1,3 @@
 from .AbstractBalancer import AbstractBalancer
 from .NoBalancer import NoBalancer
+from .ThreadBalancer import ThreadBalancer
diff --git a/QtPyNetwork/client/AbstractClient.py b/QtPyNetwork/client/AbstractClient.py
new file mode 100644
index 0000000..9db9b9d
--- /dev/null
+++ b/QtPyNetwork/client/AbstractClient.py
@@ -0,0 +1,96 @@
+from qtpy.QtCore import Slot, Signal, QObject
+
+from abc import abstractmethod
+
+
+class AbstractClient(QObject):
+
+    connected = Signal(str, int)
+    disconnected = Signal()
+    message = Signal(bytes)
+    error = Signal(Exception)
+    closed = Signal()
+    failed_to_connect = Signal()
+
+    def __init__(self, timeout: int = 5):
+        super(AbstractClient, self).__init__()
+        self._socket = None
+        self._timeout = timeout
+
+    @abstractmethod
+    @Slot(str, int)
+    def start(self, ip: str, port: int):
+        """Start client thread and connect to server."""
+        pass
+
+    @abstractmethod
+    @Slot(bytes)
+    def write(self, data: bytes):
+        """Write data to server.
+
+        Args:
+            data (bytes): Data to write.
+        """
+        pass
+
+    @Slot(str, int)
+    def on_connected(self, ip, port):
+        """Called when client connects to server.
+        Emits connected signal.
+
+        Args:
+            ip (str): Client ip address.
+            port (int): Client port.
+        """
+        self.connected.emit(ip, port)
+
+    def on_message(self, message: bytes):
+        """Called when client receives message from server.
+        Emits message signal.
+
+        Args:
+            message (bytes): Message.
+        """
+        self.message.emit(message)
+
+    @Slot()
+    def on_disconnected(self):
+        """Called when device disconnects from server.
+        Emits disconnected signal."""
+        self.disconnected.emit()
+
+    @Slot(str)
+    def on_error(self, error: str):
+        """Called when a socket error occurs.
+        Emits error signal.
+
+        Args:
+            error (str): Error string.
+        """
+        self.error.emit(Exception, error)
+
+    @Slot()
+    def on_failed_to_connect(self):
+        """Called when client fails to connect to server.
+        Emits failed_to_connect signal.
+        """
+        self.failed_to_connect.emit()
+
+    @Slot()
+    def on_closed(self):
+        """Called when the socket is closed.
+        Emits closed signal."""
+        self.closed.emit()
+
+    @Slot()
+    def close(self):
+        """Disconnect from server and close socket."""
+        if self._socket:
+            self._socket.close()
+            self._socket = None
+        self.closed.emit()
+
+    @Slot()
+    def is_running(self):
+        """Check if client is running."""
+        return self._socket is not None
diff --git a/QtPyNetwork/client/QThreadedClient.py b/QtPyNetwork/client/QThreadedClient.py
index 7204d4f..7df6ed7 100644
--- a/QtPyNetwork/client/QThreadedClient.py
+++ b/QtPyNetwork/client/QThreadedClient.py
@@ -125,8 +125,11 @@ def _disconnect(self):
 
     @Slot()
     def close(self):
+        print("Pre _disconnect")
         self._disconnect()
+        print("Pre tcpsocket.close")
         self.tcpSocket.close()
+        print("Pre closed.emit")
         self.closed.emit()
 
 
@@ -257,8 +260,11 @@ def write(self, data: bytes):
     def close(self):
         """Disconnect from server and close socket."""
         if self.__client and self.__client_thread:
+            print("__client and __client_thread")
             self.__client.close_signal.emit()
+            print("Signal emited")
             self.__client_thread.quit()
+            print("_client_thread quited")
         else:
             self.error.emit("Client not running")
 
diff --git a/QtPyNetwork/client/TCPClient.py b/QtPyNetwork/client/TCPClient.py
new file mode 100644
index 0000000..c3debfb
--- /dev/null
+++ b/QtPyNetwork/client/TCPClient.py
@@ -0,0 +1,71 @@
+from qtpy.QtNetwork import QAbstractSocket, QTcpSocket, QHostAddress
+from qtpy.QtCore import Signal, Slot, QTimer
+
+from QtPyNetwork.common import read, write
+from .AbstractClient import AbstractClient
+
+import logging
+
+
+class TCPClient(AbstractClient):
+    def __init__(self, timeout: int = 2):
+        super(TCPClient, self).__init__(timeout)
+        self._logger = logging.getLogger(self.__class__.__name__)
+        self._data = b""
+        self._size_left = 0
+
+    @Slot(str, int)
+    def start(self, ip: str, port: int):
+        self._socket = QTcpSocket(self)
+        self._socket.connected.connect(self.__on_socket_connected)
+        self._socket.disconnected.connect(self.__on_socket_disconnected)
+        self._socket.readyRead.connect(self.__on_socket_ready_read)
+        self._socket.error.connect(self.__on_socket_error)
+        self._socket.connectToHost(QHostAddress(ip), port)
+        self._logger.debug(f"Connecting to {ip}:{port}")
+        QTimer.singleShot(self._timeout * 1000, self.__check_connected)
+        self._logger.debug(f"Started connection timer with timeout {self._timeout} seconds")
+
+    @Slot(bytes)
+    def write(self, data: bytes):
+        if self._socket.state() == QAbstractSocket.ConnectedState:
+            write(self._socket, data)
+
+    @Slot()
+    def __on_socket_connected(self):
+        ip = self._socket.peerAddress().toString()
+        port = int(self._socket.peerPort())
+        self._logger.info("Connected to {}:{}".format(ip, port))
+        self.on_connected(ip, port)
+
+    @Slot()
+    def __on_socket_disconnected(self):
+        self._logger.info("Disconnected from server")
+        self.on_disconnected()
+
+    @Slot()
+    def __on_socket_ready_read(self):
+        data, size_left = read(self._socket, self._data, self._size_left)
+        print(f"TCPClient left: {size_left}, received: {data}")
+        if size_left == 0:
+            self._size_left = 0
+            self._data = b""
+            self.on_message(data)
+        else:
+            self._data = data
+            self._size_left = size_left
+
+    @Slot()
+    def __on_socket_error(self):
+        """Handle socket errors.
+
+        Note:
+            Emits error signal.
+        """
+        error = self._socket.errorString()
+        self.error.emit(Exception(error))
+
+    def __check_connected(self):
+        if not self._socket.state() == QAbstractSocket.ConnectedState:
+            self._socket.close()
+            self.on_failed_to_connect()
diff --git a/QtPyNetwork2/client/UDPClient.py b/QtPyNetwork/client/UDPClient.py
similarity index 100%
rename from QtPyNetwork2/client/UDPClient.py
rename to QtPyNetwork/client/UDPClient.py
diff --git a/QtPyNetwork/client/__init__.py b/QtPyNetwork/client/__init__.py
index 4c4296b..911f590 100644
--- a/QtPyNetwork/client/__init__.py
+++ b/QtPyNetwork/client/__init__.py
@@ -1 +1 @@
-from .QThreadedClient import QThreadedClient
+from .TCPClient import TCPClient
diff --git a/QtPyNetwork/common.py b/QtPyNetwork/common.py
new file mode 100644
index 0000000..2868681
--- /dev/null
+++ b/QtPyNetwork/common.py
@@ -0,0 +1,48 @@
+from qtpy.QtNetwork import QAbstractSocket
+from struct import unpack, calcsize, pack
+
+HEADER = '!L'
+HEADER_SIZE = calcsize(HEADER)
+
+
+def read(socket: QAbstractSocket, buffer: bytes, size_left: int) -> (bytes, int):
+    """Read data from socket.
+
+    Args:
+        socket (QAbstractSocket): Socket to read from.
+        buffer (bytes): Buffer to read into.
+        size_left (int): Size left to read.
+
+    Returns:
+        bytes: Read data.
+    """
+    while socket.bytesAvailable():
+        print(f"Reading {size_left} bytes with buffer {buffer}")
+        if size_left > 0:
+            data = socket.read(size_left)
+            size_left = size_left - len(data)
+            if size_left > 0:
+                buffer += data
+            else:
+                return buffer + data, 0
+        else:
+            header = socket.read(HEADER_SIZE)
+            data_size = unpack(HEADER, header)[0]
+            data = socket.read(data_size)
+            if len(data) < data_size:
+                buffer = data
+                size_left = data_size
+            else:
+                return data, 0
+
+
+def write(socket: QAbstractSocket, data: bytes) -> None:
+    """Write data to socket.
+
+    Args:
+        socket (QAbstractSocket): Socket to write to.
+        data (bytes): Data to write.
+    """
+    print(f"Writing {len(data)} bytes with buffer {data}")
+    data = pack(HEADER, len(data)) + data
+    socket.write(data)
diff --git a/QtPyNetwork2/models/Device.py b/QtPyNetwork/models/Client.py
similarity index 95%
rename from QtPyNetwork2/models/Device.py
rename to QtPyNetwork/models/Client.py
index 536b1ce..f882dff 100644
--- a/QtPyNetwork2/models/Device.py
+++ b/QtPyNetwork/models/Client.py
@@ -49,8 +49,8 @@ def is_connected(self) -> bool:
         return self.__connected
 
     @Slot()
-    def kick(self):
-        self.server().kick(self)
+    def disconnect(self):
+        self.server().disconnect(self)
 
     @Slot(bytes)
     def write(self, message: bytes):
diff --git a/QtPyNetwork/models/__init__.py b/QtPyNetwork/models/__init__.py
index c5b5fb3..f39db69 100644
--- a/QtPyNetwork/models/__init__.py
+++ b/QtPyNetwork/models/__init__.py
@@ -1 +1 @@
-from .Device import Device
+from .Client import Client
diff --git a/QtPyNetwork2/server/AbstractServer.py b/QtPyNetwork/server/AbstractServer.py
similarity index 94%
rename from QtPyNetwork2/server/AbstractServer.py
rename to QtPyNetwork/server/AbstractServer.py
index 7c62af8..dae5d3d 100644
--- a/QtPyNetwork2/server/AbstractServer.py
+++ b/QtPyNetwork/server/AbstractServer.py
@@ -1,12 +1,12 @@
 from qtpy.QtCore import Slot, Signal, QObject, QThread
 from qtpy.QtNetwork import QTcpServer, QHostAddress
 
-from QtPyNetwork2.models import Client
-from QtPyNetwork2.exceptions import NotConnectedError, ServerNotRunning
+from QtPyNetwork.models import Client
+from QtPyNetwork.exceptions import NotConnectedError, ServerNotRunning
 
 import logging
 
-from QtPyNetwork2.balancers import AbstractBalancer
+from QtPyNetwork.balancer import AbstractBalancer
 
 
 class AbstractServer(QObject):
@@ -24,7 +24,7 @@ def __init__(self, balancer: AbstractBalancer):
         super(AbstractServer, self).__init__()
         self.clients: list[Client] = []
         self.server: QObject = None
-        self.client_model = Client
+        self.__client_model = Client
 
         self.balancer = balancer
         self.balancer.connected.connect(self.__on_balancer_client_connected)
@@ -35,7 +35,7 @@ def __init__(self, balancer: AbstractBalancer):
 
     @Slot(int, str, int)
     def __on_balancer_client_connected(self, client_id: int, ip: str, port: int):
-        client = self.client_model(self, client_id, ip, port)
+        client = self.__client_model(self, client_id, ip, port)
         self.clients.append(client)
         self.on_connected(client, ip, port)
 
@@ -48,10 +48,10 @@ def __on_balancer_client_message(self, client_id: int, message: bytes):
     def __on_balancer_client_disconnected(self, client_id: int):
         """When client disconnects from server."""
         client = self.get_client_by_id(client_id)
-        client.set_connected(False)
-        if client in self.clients:
+        if client:
+            client.set_connected(False)
             self.clients.remove(client)
-        self.on_disconnected(client)
+            self.on_disconnected(client)
 
     @Slot(int, Exception)
     def __on_balancer_client_error(self, client_id: int, error: Exception):
@@ -123,6 +123,15 @@ def on_server_error(self, error: Exception):
     def on_closed(self):
         self.closed.emit()
 
+    @Slot(Client)
+    def disconnect(self, client: Client):
+        """Disconnects client from server.
+
+        Args:
+            client (Client): Client object.
+        """
+        self.balancer.disconnect(client.id())
+
     @Slot(Client, bytes)
     def write(self, client: Client, message: bytes):
         """Sends message to client.
@@ -148,6 +157,12 @@ def get_client_by_id(self, client_id: int):
             if client.id() == client_id:
                 return client
 
+    @Slot(Client)
+    def set_client_model(self, model: Client):
+        if not issubclass(model, Client):
+            raise TypeError('model must be subclass of Client')
+        self.__client_model = model
+
     @Slot()
     def is_running(self) -> bool:
         """Check if server is running."""
diff --git a/QtPyNetwork2/server/TCPServer.py b/QtPyNetwork/server/TCPServer.py
similarity index 100%
rename from QtPyNetwork2/server/TCPServer.py
rename to QtPyNetwork/server/TCPServer.py
diff --git a/QtPyNetwork2/server/UDPServer.py b/QtPyNetwork/server/UDPServer.py
similarity index 100%
rename from QtPyNetwork2/server/UDPServer.py
rename to QtPyNetwork/server/UDPServer.py
diff --git a/QtPyNetwork/server/__init__.py b/QtPyNetwork/server/__init__.py
index 5e60b9a..f91113b 100644
--- a/QtPyNetwork/server/__init__.py
+++ b/QtPyNetwork/server/__init__.py
@@ -1,2 +1 @@
-from .QBalancedServer import QBalancedServer
-from .QThreadedServer import QThreadedServer
+from .TCPServer import TCPServer
diff --git a/QtPyNetwork2/client/AbstractClient.py b/QtPyNetwork2/client/AbstractClient.py
deleted file mode 100644
index e69de29..0000000
diff --git a/QtPyNetwork2/client/TCPClient.py b/QtPyNetwork2/client/TCPClient.py
deleted file mode 100644
index e69de29..0000000
diff --git a/QtPyNetwork2/models/__init__.py b/QtPyNetwork2/models/__init__.py
deleted file mode 100644
index e114aa1..0000000
--- a/QtPyNetwork2/models/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from .Device import Client, Device
diff --git a/QtPyNetwork2/server/__init__.py b/QtPyNetwork2/server/__init__.py
deleted file mode 100644
index f91113b..0000000
--- a/QtPyNetwork2/server/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from .TCPServer import TCPServer
diff --git a/QtPyNetwork2/__init__.py b/examples/QtPyNetwork/__init__.py
similarity index 100%
rename from QtPyNetwork2/__init__.py
rename to examples/QtPyNetwork/__init__.py
diff --git a/QtPyNetwork2/__main__.py b/examples/QtPyNetwork/__main__.py
similarity index 100%
rename from QtPyNetwork2/__main__.py
rename to examples/QtPyNetwork/__main__.py
diff --git a/QtPyNetwork2/client/QThreadedClient.py b/examples/QtPyNetwork/client/QThreadedClient.py
similarity index 97%
rename from QtPyNetwork2/client/QThreadedClient.py
rename to examples/QtPyNetwork/client/QThreadedClient.py
index 7df6ed7..7204d4f 100644
--- a/QtPyNetwork2/client/QThreadedClient.py
+++ b/examples/QtPyNetwork/client/QThreadedClient.py
@@ -125,11 +125,8 @@ def _disconnect(self):
 
     @Slot()
     def close(self):
-        print("Pre _disconnect")
         self._disconnect()
-        print("Pre tcpsocket.close")
         self.tcpSocket.close()
-        print("Pre closed.emit")
         self.closed.emit()
 
 
@@ -260,11 +257,8 @@ def write(self, data: bytes):
     def close(self):
         """Disconnect from server and close socket."""
         if self.__client and self.__client_thread:
-            print("__client and __client_thread")
             self.__client.close_signal.emit()
-            print("Signal emited")
             self.__client_thread.quit()
-            print("_client_thread quited")
         else:
             self.error.emit("Client not running")
 
diff --git a/QtPyNetwork2/client/__init__.py b/examples/QtPyNetwork/client/__init__.py
similarity index 100%
rename from QtPyNetwork2/client/__init__.py
rename to examples/QtPyNetwork/client/__init__.py
diff --git a/QtPyNetwork2/exceptions/__init__.py b/examples/QtPyNetwork/exceptions/__init__.py
similarity index 100%
rename from QtPyNetwork2/exceptions/__init__.py
rename to examples/QtPyNetwork/exceptions/__init__.py
diff --git a/QtPyNetwork/models/Device.py b/examples/QtPyNetwork/models/Device.py
similarity index 100%
rename from QtPyNetwork/models/Device.py
rename to examples/QtPyNetwork/models/Device.py
diff --git a/examples/QtPyNetwork/models/__init__.py b/examples/QtPyNetwork/models/__init__.py
new file mode 100644
index 0000000..c5b5fb3
--- /dev/null
+++ b/examples/QtPyNetwork/models/__init__.py
@@ -0,0 +1 @@
+from .Device import Device
diff --git a/QtPyNetwork/server/BaseServer.py b/examples/QtPyNetwork/server/BaseServer.py
similarity index 99%
rename from QtPyNetwork/server/BaseServer.py
rename to examples/QtPyNetwork/server/BaseServer.py
index 7eb79f3..fee80fe 100644
--- a/QtPyNetwork/server/BaseServer.py
+++ b/examples/QtPyNetwork/server/BaseServer.py
@@ -166,7 +166,7 @@ def write(self, device: Device, data: bytes):
 
     @Slot(bytes)
     def write_all(self, data: bytes):
-        """Write data to all devices."""
+        """Write data to all clients."""
         if not self.__server or not self.__handler:
             raise ServerNotRunning("Server is not running")
         self.__handler.write_all.emit(data)
@@ -234,7 +234,7 @@ def get_device_by_id(self, device_id: int) -> Device:
         raise Exception("CLIENT-{} not found".format(device_id))
 
     def get_devices(self):
-        """Returns list with devices."""
+        """Returns list with clients."""
         return self.__devices
 
     def set_handler_class(self, handler):
diff --git a/QtPyNetwork/server/QBalancedServer.py b/examples/QtPyNetwork/server/QBalancedServer.py
similarity index 100%
rename from QtPyNetwork/server/QBalancedServer.py
rename to examples/QtPyNetwork/server/QBalancedServer.py
diff --git a/QtPyNetwork/server/QThreadedServer.py b/examples/QtPyNetwork/server/QThreadedServer.py
similarity index 100%
rename from QtPyNetwork/server/QThreadedServer.py
rename to examples/QtPyNetwork/server/QThreadedServer.py
diff --git a/examples/QtPyNetwork/server/__init__.py b/examples/QtPyNetwork/server/__init__.py
new file mode 100644
index 0000000..5e60b9a
--- /dev/null
+++ b/examples/QtPyNetwork/server/__init__.py
@@ -0,0 +1,2 @@
+from .QBalancedServer import QBalancedServer
+from .QThreadedServer import QThreadedServer
diff --git a/examples2/balanced_server.py b/examples2/balanced_server.py
index 30a1d8a..f5e6d10 100644
--- a/examples2/balanced_server.py
+++ b/examples2/balanced_server.py
@@ -5,8 +5,8 @@
 import logging
 import traceback
 
-from QtPyNetwork2.server import TCPServer
-from QtPyNetwork2.balancers import ThreadPoolBalancer
+from QtPyNetwork.server import TCPServer
+from QtPyNetwork.balancer import ThreadPoolBalancer
 from QtPyNetwork.models import Device
 
 IP = "127.0.0.1"
diff --git a/examples2/client.py b/examples2/client.py
index 6d41eb1..b2b6ccc 100644
--- a/examples2/client.py
+++ b/examples2/client.py
@@ -1,10 +1,11 @@
 from qtpy.QtWidgets import QApplication
-from qtpy.QtCore import QObject, Slot, QCoreApplication
+from qtpy.QtCore import QObject, Slot, QCoreApplication, qInstallMessageHandler
 
 import sys
 import logging
 
-from QtPyNetwork.client import QThreadedClient
+from QtPyNetwork.client import TCPClient
+from qrainbowstyle.extras import qt_message_handler
 
 IP = "127.0.0.1"
 PORT = 12500
@@ -14,31 +15,31 @@ class Main(QObject):
 
     def __init__(self):
         super(Main, self).__init__(None)
-
-    def setup(self):
         self.logger = logging.getLogger(self.__class__.__name__)
 
-        self.cln = QThreadedClient()
-        self.cln.message.connect(self.on_message)
-        self.cln.connected.connect(self.on_connected)
-        self.cln.failed_to_connect.connect(self.close)
-        self.cln.disconnected.connect(self.close)
-        self.cln.start(IP, PORT)
+        self.client = TCPClient()
+        self.client.message.connect(self.on_message)
+        self.client.connected.connect(self.on_connected)
+        self.client.failed_to_connect.connect(self.on_failed_to_connect)
+        self.client.disconnected.connect(self.close)
+        self.client.start(IP, PORT)
 
     @Slot(str, int)
     def on_connected(self, ip: str, port: int):
         self.logger.info(f"Connected to {ip}:{port}")
-        self.cln.write(b"Kick me plz")
+        self.client.write(b"Kick me plz")
 
     @Slot(bytes)
     def on_message(self, data: bytes):
         self.logger.info(f"Received: {data}")
 
+    @Slot()
+    def on_failed_to_connect(self):
+        self.logger.error("Failed to connect")
+
     @Slot()
     def close(self):
-        self.cln.close()
-        while self.cln.is_running():
-            self.cln.wait()
+        self.client.close()
         QApplication.instance().quit()
 
 
@@ -49,8 +50,8 @@ def close(self):
         handlers=[logging.StreamHandler()])
     logging.getLogger().debug("Logger enabled")
 
+    qInstallMessageHandler(qt_message_handler)
     app = QCoreApplication(sys.argv)
 
     main = Main()
-    main.setup()
     sys.exit(app.exec_())
diff --git a/examples2/no_balancer.py b/examples2/no_balancer.py
index 3be3c2d..e8d493d 100644
--- a/examples2/no_balancer.py
+++ b/examples2/no_balancer.py
@@ -5,9 +5,9 @@
 import logging
 import traceback
 
-from QtPyNetwork2.server import TCPServer
-from QtPyNetwork2.balancers import NoBalancer
-from QtPyNetwork2.models import Client
+from QtPyNetwork.server import TCPServer
+from QtPyNetwork.balancer import NoBalancer
+from QtPyNetwork.models import Client
 
 IP = "127.0.0.1"
 PORT = 12500
diff --git a/examples2/threaded_server.py b/examples2/threaded_server.py
index e8781e1..efa7993 100644
--- a/examples2/threaded_server.py
+++ b/examples2/threaded_server.py
@@ -5,8 +5,9 @@
 import logging
 import traceback
 
-from QtPyNetwork.server import QThreadedServer
-from QtPyNetwork.models import Device
+from QtPyNetwork.server import TCPServer
+from QtPyNetwork.balancer import ThreadBalancer
+from QtPyNetwork.models import Client
 
 IP = "127.0.0.1"
 PORT = 12500
@@ -16,32 +17,31 @@ class Main(QObject):
 
     def __init__(self):
         super(Main, self).__init__(None)
-
-    def setup(self):
         self.logger = logging.getLogger(self.__class__.__name__)
 
-        self.srv = QThreadedServer()
-        self.srv.connected.connect(lambda device, ip, port: device.write(b"Some important data"))
-        self.srv.disconnected.connect(self.on_disconnected)
-        self.srv.message.connect(self.on_message)
-        self.srv.start(IP, PORT)
+        self.server = TCPServer(ThreadBalancer())
+        self.server.connected.connect(lambda device, ip, port: device.write(b"Some important data"))
+        self.server.disconnected.connect(self.on_disconnected)
+        self.server.message.connect(self.on_message)
+        self.server.start(IP, PORT)
 
-    @Slot(Device, bytes)
-    def on_message(self, device, message: bytes):
-        self.logger.info("Received {}: {}".format(device.id(), message))
+    @Slot(Client, bytes)
+    def on_message(self, client: Client, message: bytes):
+        self.logger.info("Received {}: {}".format(client.id(), message))
         if message.decode() == "Kick me plz":
-            device.kick()
+            client.write(b"I'm kicked")
+            client.disconnect()
 
-    @Slot(Device)
+    @Slot(Client)
     def on_disconnected(self, device):
         self.logger.info("Disconnected: {}; Connected: {}".format(device.id(), device.is_connected()))
-        self.close()
+        # self.close()
 
     @Slot()
     def close(self):
-        self.srv.close()
-        while self.srv.is_running():
-            self.srv.wait()
+        self.server.close()
+        while self.server.is_running():
+            self.server.wait()
         QApplication.instance().quit()
 
 
@@ -62,5 +62,4 @@ def exception_hook(exctype, value, tb):
     app = QCoreApplication(sys.argv)
 
     main = Main()
-    main.setup()
     sys.exit(app.exec_())