Fix the expand button changing too soon

The expand button used to change its icon right after it was clicked and
before the animation has finished. This resulted in the icon being changed too soon. This CL fixes this by waiting until the pip transition is finished

Bug: 237269877

Test: Manual
Change-Id: I46d8bc82be7db0532c51aa116ff2dbbf5e09fc68