refactor tests slightly