Remove SetFlagsRule from DocsUI tests and replace them with CheckFlagsRule.

SetFlagsRule is not supported in mts. Instead we should filter tests to
test the code under both states of the flags.

I couldn't reproduce the mts issue locally so not certain this fixes it.

Bug: 390031738
Test: Fixing unit tests. No production code change.
Flag: EXEMPT test only
Change-Id: I3f46bc160eb8819dd4d35a8f942480d2b4c9aa92