Improve UI tests