Refactor and add start of some tests