Increase test coverage and fix bugs exposed