Added a bit more logic/tests stuff