Adjusting tests to match new code structure.