Fixed bug in test suite
Fixed critical bug + minor refactoring