Unit tests bugfix