clean up a little
Add some unit tests