UI changes and test fixes