[WIP] Cleaning up and trying to fix test failure

See #20 #18