Fixed some issues arised with tests after merge