[expectations] Refactor code