Tidy up following PR review