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.