Code cleanup
Update tests