Minor changes to find elements and try to fix broken test