Code review clean up