code cleanup, tests passing