More code cleanup.  Added some tests..