This commit replaces `waitForFunction` calls that use `document.activeElement` to wait for an element to get focus by simpler `waitForSelector` expressions that use the `:focus` selector. Note that we already use this in other tests, so this improves consistency too.