Code cleanup
Add test for #159