Bug fixes + additional tests