Some fixes + tests