Few timeout fixes related to UI

Increase the timeout value to avoid timing issues.

Used `wait_until_element` as the element takes time to appear.
Uhhhhh wow. Yeah.