Adding test, refactors and fixing connection graph