Fix UI unit tests