Massive clean up and writing tests