Refactored connection slightly and started to write tests for it.