Add tests, add error class