Added tests and a few bug fixes