Some fixes and added some tests