Code cleanup
Harden api and add tests