Cleanups and unit test for #166