More refactoring of button handling