Improved tests and did some formatting/refactoring

TODO: Controller tests