Fix unittests

Fix most of the failing unit tests. Should fix some of the failures in #118