Cleaned up code and added more structure
test