Updated tests to include latest schema changes