Address review comments on tests / generators