Fix shit, add tests