Fixed tests to go with new structure