Minor fix in tests