Adding test coverage (WIP)