Minor tweaks to ensure test mode operates as it should