Cleanup dead code and add further tests