Change logic of test