Fixed bug after refactor
Cleaned up code and added tests