add if statement for tests