Work on tests! Fixed some issues :)