Czat napisany uzywajacy socketow TCP zawiesza sie

Po krotkim czasie bezczynosci moj czat tak jakby zawiesza sie. Nie ma zerwania polaczenia i nie mozna wyslac zadnej wiadomosci. Uzywam protokolu TCP

import socket 
import sys 

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

server_adress = ('localhost',8081)

sock.bind(server_adress)

sock.listen(1)

while True:
    print('waiting for a connection')
    connection, client_address = sock.accept()
    try:
        print('connection from', client_address)
        while True:
            data = connection.recv(16)
            print(str(data, "utf8"))
            if data:
                message = bytes(input(), 'utf-8')
                connection.sendall(message)
            else:
                print('no data from', client_address)
                break

    finally:
        connection.close()

import socket
import sys


sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

server_address = ('localhost', 8081)
print('connecting to {} port {}'.format(*server_address))
sock.connect(server_address)
while True:

	str_message = input()
	byte_message = bytes( str_message, 'utf-8')
	if str_message=='end':
		print('closing socket')
		sock.close()
		break
	sock.sendall(byte_message)

	amount_received = 0
	amount_expected = len(byte_message)

	while amount_received < amount_expected:
		data = sock.recv(16)
		amount_received += len(data)
		print(str(data, "utf8"))

Spróbuj ustawić niezerowy timeout lub non-blocking mode, ale najpierw postaw breakpointy w debuggerze na wysyłaniu/odbieraniu i sprawdź co dokładnie blokuje ;)