Fix tests

* Basic feature for UI.