Updated first unit test to work with current architecture