Bugfix + added tests