Fix socket connection and add tests