Working unit tests and bugfixes