Clean up test environment and fix tests