[#008] Fixing last bugs introduced with refactoring