Bug fixes + tests