EspressoでViewのVisibilityが変化するまで待つようにしてみた

ボタンを押したらあるViewのVisibilityをAnimationさせて変化させている。
AndroidのEspressoでそこの部分をUIテストしようとすると、ボタンを押してからVisibilityが変化したことを確認するまでにThread.sleepで少し待たなければいけない。
しかし、必ず何秒か待つのはスマートではない。

android - Espresso: Thread.sleep( ); - Stack Overflowの記事を見るともうちょっと良いやり方があったので参考にして書いてみた。

以下のように書くと対象のViewのVisibilityがVISIBLEになるまで最大1000ミリ秒待つようにしてみた。

onView(withId(R.id.controller)).perform(waitVisibility(View.VISIBLE, 1000));

waitVisibility関数の中身

    public static ViewAction waitVisibility(int visibility, long millis) {
        return new WaitVisibility(visibility, millis);
    }

    public static abstract class WaitBase implements ViewAction {
        protected long millis = 0;

        WaitBase(long millis) {
            this.millis = millis;
        }

        @Override
        public Matcher<View> getConstraints() {
            return new IsAnything<>();
        }

        @Override
        public String getDescription() {
            return"wait during " + millis + " millis.";
        }

        @Override
        public void perform(UiController uiController, View view) {
            uiController.loopMainThreadUntilIdle();
            final long startTime = System.currentTimeMillis();
            final long endTime = startTime + millis;

            do {
                if (check(view)) {
                    return;
                }

                uiController.loopMainThreadForAtLeast(50);
            }
            while (System.currentTimeMillis() < endTime);

            // timeout happens
            throw new PerformException.Builder()
                    .withActionDescription(this.getDescription())
                    .withViewDescription(HumanReadables.describe(view))
                    .withCause(new TimeoutException())
                    .build();
        }


        abstract protected boolean check(final View view);
    }

    public static class WaitVisibility extends WaitBase {
        private final int mVisibility;

        public WaitVisibility(int visibility, long millis) {
            super(millis);
            mVisibility = visibility;
        }

        @Override
        public String getDescription() {
            return "wait for a specific view with visibility <" + mVisibility + "> during " + millis + " millis.";
        }

        @Override
        protected boolean check(final View view) {
            return view.getVisibility() == mVisibility;
        }
    }