Additional changes to fix bugs and update tests.