[wip] Work on cleanup and tests