Making more testable and adding tests