Tweaked unit tests to reflect new structure