Updated tests for refactored schema