Code cleanup
Expand tests