Post code review fixes