[DF-2836] Fix: refactor due to review