diff --git a/camera/camera-camera2/src/test/java/androidx/camera/camera2/internal/TorchControlTest.java b/camera/camera-camera2/src/test/java/androidx/camera/camera2/internal/TorchControlTest.java
index 5d2a7be..b60734c 100644
--- a/camera/camera-camera2/src/test/java/androidx/camera/camera2/internal/TorchControlTest.java
+++ b/camera/camera-camera2/src/test/java/androidx/camera/camera2/internal/TorchControlTest.java
@@ -16,6 +16,8 @@
 
 package androidx.camera.camera2.internal;
 
+import static android.os.Looper.getMainLooper;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.Mockito.mock;
@@ -23,6 +25,7 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.mockito.internal.verification.VerificationModeFactory.times;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.content.Context;
 import android.hardware.camera2.CameraAccessException;
@@ -123,6 +126,7 @@
     public void enableTorch_whenNoFlashUnit() throws InterruptedException {
         Throwable cause = null;
         try {
+            // Without a flash unit, this future will complete immediately. No need to idle.
             mNoFlashUnitTorchControl.enableTorch(true).get();
         } catch (ExecutionException e) {
             // The real cause is wrapped in ExecutionException, retrieve it and check.
@@ -133,16 +137,22 @@
 
     @Test
     public void getTorchState_whenNoFlashUnit() {
-        int torchState = mNoFlashUnitTorchControl.getTorchState().getValue();
+        int torchState =
+                Objects.requireNonNull(mNoFlashUnitTorchControl.getTorchState().getValue());
         assertThat(torchState).isEqualTo(TorchState.OFF);
     }
 
     @Test
     public void enableTorch_whenInactive() throws InterruptedException {
         mTorchControl.setActive(false);
+        ListenableFuture<Void> listenableFuture = mTorchControl.enableTorch(true);
+        // enableTorch can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
+
+        assertThat(listenableFuture.isDone()).isTrue();
         Throwable cause = null;
         try {
-            mTorchControl.enableTorch(true).get();
+            listenableFuture.get();
         } catch (ExecutionException e) {
             // The real cause is wrapped in ExecutionException, retrieve it and check.
             cause = e.getCause();
@@ -154,7 +164,8 @@
     @Test
     public void getTorchState_whenInactive() {
         mTorchControl.setActive(false);
-        int torchState = mTorchControl.getTorchState().getValue();
+        // LiveData is updated synchronously. No need to idle.
+        int torchState = Objects.requireNonNull(mTorchControl.getTorchState().getValue());
 
         assertThat(torchState).isEqualTo(TorchState.OFF);
     }
@@ -162,7 +173,8 @@
     @Test
     public void enableTorch_torchStateOn() {
         mTorchControl.enableTorch(true);
-        int torchState = mTorchControl.getTorchState().getValue();
+        // LiveData is updated synchronously. No need to idle.
+        int torchState = Objects.requireNonNull(mTorchControl.getTorchState().getValue());
 
         assertThat(torchState).isEqualTo(TorchState.ON);
     }
@@ -170,26 +182,34 @@
     @Test
     public void disableTorch_TorchStateOff() {
         mTorchControl.enableTorch(true);
-        int torchState = mTorchControl.getTorchState().getValue();
-
-        assertThat(torchState).isEqualTo(TorchState.ON);
+        // LiveData is updated synchronously. No need to idle.
+        int firstTorchState = Objects.requireNonNull(mTorchControl.getTorchState().getValue());
 
         mTorchControl.enableTorch(false);
-        torchState = mTorchControl.getTorchState().getValue();
+        // LiveData is updated synchronously. No need to idle.
+        int secondTorchState = mTorchControl.getTorchState().getValue();
 
-        assertThat(torchState).isEqualTo(TorchState.OFF);
+        assertThat(firstTorchState).isEqualTo(TorchState.ON);
+        assertThat(secondTorchState).isEqualTo(TorchState.OFF);
     }
 
-    @Test(timeout = 5000L)
+    @Test
     public void enableDisableTorch_futureWillCompleteSuccessfully()
             throws ExecutionException, InterruptedException {
         ListenableFuture<Void> future = mTorchControl.enableTorch(true);
+        // enableTorch can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
+
+        // Calling onCaptureResult directly from executor thread (main thread). No need to idle.
         mCaptureResultListener.onCaptureResult(
                 mockFlashCaptureResult(CaptureResult.FLASH_MODE_TORCH));
         // Future should return with no exception
         future.get();
 
         future = mTorchControl.enableTorch(false);
+        // enableTorch can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
+        // Calling onCaptureResult directly from executor thread (main thread). No need to idle.
         mCaptureResultListener.onCaptureResult(
                 mockFlashCaptureResult(CaptureResult.FLASH_MODE_OFF));
         // Future should return with no exception
@@ -200,6 +220,8 @@
     public void enableTorchTwice_cancelPreviousFuture() throws InterruptedException {
         ListenableFuture<Void> future = mTorchControl.enableTorch(true);
         mTorchControl.enableTorch(true);
+        // enableTorch can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
         Throwable cause = null;
         try {
             future.get();
@@ -214,6 +236,9 @@
     @Test
     public void setInActive_cancelPreviousFuture() throws InterruptedException {
         ListenableFuture<Void> future = mTorchControl.enableTorch(true);
+        // enableTorch can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         mTorchControl.setActive(false);
         Throwable cause = null;
         try {
@@ -229,20 +254,24 @@
     @Test
     public void setInActiveWhenTorchOn_changeToTorchOff() {
         mTorchControl.enableTorch(true);
-        int torchState = mTorchControl.getTorchState().getValue();
+        // enableTorch can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
+        int initialTorchState = Objects.requireNonNull(mTorchControl.getTorchState().getValue());
 
-        assertThat(torchState).isEqualTo(TorchState.ON);
-
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         mTorchControl.setActive(false);
-        torchState = mTorchControl.getTorchState().getValue();
+        int torchStateAfterInactive = mTorchControl.getTorchState().getValue();
 
-        assertThat(torchState).isEqualTo(TorchState.OFF);
+        assertThat(initialTorchState).isEqualTo(TorchState.ON);
+        assertThat(torchStateAfterInactive).isEqualTo(TorchState.OFF);
     }
 
     @Test
     public void enableDisableTorch_observeTorchStateLiveData() {
+        @SuppressWarnings("unchecked")
         Observer<Integer> observer = mock(Observer.class);
         LiveData<Integer> torchStateLiveData = mTorchControl.getTorchState();
+        // Adding observer from main thread should synchronously be notified of initial state
         torchStateLiveData.observe(mLifecycleOwner, new Observer<Integer>() {
             private Integer mValue;
             @Override
@@ -254,7 +283,12 @@
         });
 
         mTorchControl.enableTorch(true);
+        // Idle the main thread to receive first update
+        shadowOf(getMainLooper()).idle();
+
         mTorchControl.enableTorch(false);
+        // Idle the main thread to receive second update
+        shadowOf(getMainLooper()).idle();
 
         ArgumentCaptor<Integer> torchStateCaptor = ArgumentCaptor.forClass(Integer.class);
         verify(observer, times(3)).onChanged(torchStateCaptor.capture());
diff --git a/camera/camera-camera2/src/test/java/androidx/camera/camera2/internal/ZoomControlTest.java b/camera/camera-camera2/src/test/java/androidx/camera/camera2/internal/ZoomControlTest.java
index 94ae05c..f6cce84 100644
--- a/camera/camera-camera2/src/test/java/androidx/camera/camera2/internal/ZoomControlTest.java
+++ b/camera/camera-camera2/src/test/java/androidx/camera/camera2/internal/ZoomControlTest.java
@@ -16,14 +16,15 @@
 
 package androidx.camera.camera2.internal;
 
+import static android.os.Looper.getMainLooper;
+
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.content.Context;
 import android.graphics.Rect;
@@ -35,12 +36,9 @@
 import android.os.Build;
 
 import androidx.annotation.NonNull;
-import androidx.annotation.Nullable;
 import androidx.camera.core.CameraControl;
 import androidx.camera.core.impl.CameraControlInternal;
 import androidx.camera.core.impl.utils.executor.CameraXExecutors;
-import androidx.camera.core.impl.utils.futures.FutureCallback;
-import androidx.camera.core.impl.utils.futures.Futures;
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.filters.SmallTest;
 
@@ -57,9 +55,8 @@
 import org.robolectric.shadows.ShadowCameraCharacteristics;
 import org.robolectric.shadows.ShadowCameraManager;
 
-import java.util.concurrent.CountDownLatch;
+import java.util.Objects;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
 
 @SmallTest
 @RunWith(RobolectricTestRunner.class)
@@ -76,6 +73,10 @@
     private ZoomControl mZoomControl;
     private Camera2CameraControl.CaptureResultListener mCaptureResultListener;
 
+    private static Rect getCropRectByRatio(float ratio) {
+        return ZoomControl.getCropRectByRatio(SENSOR_RECT, ratio);
+    }
+
     @Before
     public void setUp() throws CameraAccessException {
         initShadowCameraManager();
@@ -133,32 +134,20 @@
                 .addCamera(CAMERA1_ID, characteristics1);
     }
 
-    private static Rect getCropRectByRatio(float ratio) {
-        return ZoomControl.getCropRectByRatio(SENSOR_RECT, ratio);
-    }
-
     @Test
     public void setZoomRatio1_whenResultCropRegionIsAlive_ListenableFutureSucceeded()
-            throws InterruptedException {
-        CountDownLatch latch = new CountDownLatch(1);
+            throws InterruptedException, ExecutionException {
         ListenableFuture<Void> listenableFuture = mZoomControl.setZoomRatio(1.0f);
+        // setZoomRatio can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
 
         TotalCaptureResult result = mockCaptureResult(getCropRectByRatio(1.0f));
+        // Calling onCaptureResult directly from executor thread (main thread). No need to idle.
         mCaptureResultListener.onCaptureResult(result);
 
-        Futures.addCallback(listenableFuture, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-                latch.countDown();
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
-        assertTrue(latch.await(500, TimeUnit.MILLISECONDS));
+        assertThat(listenableFuture.isDone()).isTrue();
+        // Future should have succeeded. Should not throw.
+        listenableFuture.get();
     }
 
     @NonNull
@@ -172,295 +161,227 @@
 
     @Test
     public void setZoomRatioOtherThan1_whenResultCropRegionIsAlive_ListenableFutureSucceeded()
-            throws InterruptedException {
-        CountDownLatch latch = new CountDownLatch(1);
+            throws InterruptedException, ExecutionException {
         ListenableFuture<Void> listenableFuture = mZoomControl.setZoomRatio(2.0f);
+        // setZoomRatio can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
 
         TotalCaptureResult result2 = mockCaptureResult(getCropRectByRatio(2.0f));
+        // Calling onCaptureResult directly from executor thread (main thread). No need to idle.
         mCaptureResultListener.onCaptureResult(result2);
 
-        Futures.addCallback(listenableFuture, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-                latch.countDown();
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
-        assertTrue(latch.await(500, TimeUnit.MILLISECONDS));
+        assertThat(listenableFuture.isDone()).isTrue();
+        // Future should have succeeded. Should not throw.
+        listenableFuture.get();
     }
 
     @Test
     public void setLinearZoom_valueIsAlive_ListenableFutureSucceeded()
-            throws InterruptedException {
-        CountDownLatch latch = new CountDownLatch(1);
+            throws ExecutionException, InterruptedException {
         ListenableFuture<Void> listenableFuture = mZoomControl.setLinearZoom(0.1f);
+        // setLinearZoom can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
 
-        float targetRatio = mZoomControl.getZoomState().getValue().getZoomRatio();
+        float targetRatio = Objects.requireNonNull(
+                mZoomControl.getZoomState().getValue()).getZoomRatio();
         TotalCaptureResult result = mockCaptureResult(getCropRectByRatio(targetRatio));
+        // Calling onCaptureResult directly from executor thread (main thread). No need to idle.
         mCaptureResultListener.onCaptureResult(result);
 
-        Futures.addCallback(listenableFuture, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-                latch.countDown();
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
-        assertTrue(latch.await(500, TimeUnit.MILLISECONDS));
+        assertThat(listenableFuture.isDone()).isTrue();
+        // Future should have succeeded. Should not throw.
+        listenableFuture.get();
     }
 
     @Test
-    public void setZoomRatio_newRatioIsSet_operationCanceled() throws InterruptedException {
-        CountDownLatch latchForOp1Canceled = new CountDownLatch(1);
-        CountDownLatch latchForOp2Canceled = new CountDownLatch(1);
-        CountDownLatch latchForOp3Succeeded = new CountDownLatch(1);
+    public void setZoomRatio_newRatioIsSet_operationCanceled()
+            throws InterruptedException, ExecutionException {
         ListenableFuture<Void> listenableFuture = mZoomControl.setZoomRatio(2.0f);
         ListenableFuture<Void> listenableFuture2 = mZoomControl.setZoomRatio(3.0f);
         ListenableFuture<Void> listenableFuture3 = mZoomControl.setZoomRatio(4.0f);
+        // setZoomRatio can be called from any thread and posts to executor, so idle our executor.
+        // Since multiple calls to setZoomRatio are posted in order, we only need to idle once to
+        // run all of them.
+        shadowOf(getMainLooper()).idle();
 
         TotalCaptureResult result = mockCaptureResult(getCropRectByRatio(4.0f));
+        // Calling onCaptureResult directly from executor thread (main thread). No need to idle.
         mCaptureResultListener.onCaptureResult(result);
 
-        Futures.addCallback(listenableFuture, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-            }
+        assertThat(listenableFuture.isDone()).isTrue();
+        assertThat(listenableFuture2.isDone()).isTrue();
+        assertThat(listenableFuture3.isDone()).isTrue();
+        // Futures 1 and 2 should have failed.
+        Throwable t = null;
+        try {
+            listenableFuture.get();
+        } catch (ExecutionException e) {
+            t = e.getCause();
+        }
+        assertThat(t).isInstanceOf(CameraControl.OperationCanceledException.class);
 
-            @Override
-            public void onFailure(Throwable t) {
-                if (t instanceof CameraControl.OperationCanceledException) {
-                    latchForOp1Canceled.countDown();
-                }
-            }
-        }, CameraXExecutors.mainThreadExecutor());
+        t = null;
+        try {
+            listenableFuture2.get();
+        } catch (ExecutionException e) {
+            t = e.getCause();
+        }
+        assertThat(t).isInstanceOf(CameraControl.OperationCanceledException.class);
 
-        Futures.addCallback(listenableFuture2, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-                if (t instanceof CameraControl.OperationCanceledException) {
-                    latchForOp2Canceled.countDown();
-                }
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
-        Futures.addCallback(listenableFuture3, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-                latchForOp3Succeeded.countDown();
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
-        assertTrue(latchForOp1Canceled.await(500, TimeUnit.MILLISECONDS));
-        assertTrue(latchForOp2Canceled.await(500, TimeUnit.MILLISECONDS));
-        assertTrue(latchForOp3Succeeded.await(500, TimeUnit.MILLISECONDS));
+        // Future 3 should have succeeded. Should not throw.
+        listenableFuture3.get();
     }
 
     @Test
     public void setLinearZoom_newPercentageIsSet_operationCanceled()
-            throws InterruptedException {
-        CountDownLatch latchForOp1Canceled = new CountDownLatch(1);
-        CountDownLatch latchForOp2Canceled = new CountDownLatch(1);
-        CountDownLatch latchForOp3Succeeded = new CountDownLatch(1);
+            throws InterruptedException, ExecutionException {
         ListenableFuture<Void> listenableFuture = mZoomControl.setLinearZoom(0.1f);
         ListenableFuture<Void> listenableFuture2 = mZoomControl.setLinearZoom(0.2f);
         ListenableFuture<Void> listenableFuture3 = mZoomControl.setLinearZoom(0.3f);
-        float ratioForPercentage = mZoomControl.getZoomState().getValue().getZoomRatio();
+        // setZoomRatio can be called from any thread and posts to executor, so idle our executor.
+        // Since multiple calls to setZoomRatio are posted in order, we only need to idle once to
+        // run all of them.
+        shadowOf(getMainLooper()).idle();
+        float ratioForPercentage = Objects.requireNonNull(
+                mZoomControl.getZoomState().getValue()).getZoomRatio();
 
         TotalCaptureResult result = mockCaptureResult(getCropRectByRatio(ratioForPercentage));
+        // Calling onCaptureResult directly from executor thread (main thread). No need to idle.
         mCaptureResultListener.onCaptureResult(result);
 
-        Futures.addCallback(listenableFuture, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-            }
+        assertThat(listenableFuture.isDone()).isTrue();
+        assertThat(listenableFuture2.isDone()).isTrue();
+        assertThat(listenableFuture3.isDone()).isTrue();
+        // Futures 1 and 2 should have failed.
+        Throwable t = null;
+        try {
+            listenableFuture.get();
+        } catch (ExecutionException e) {
+            t = e.getCause();
+        }
+        assertThat(t).isInstanceOf(CameraControl.OperationCanceledException.class);
 
-            @Override
-            public void onFailure(Throwable t) {
-                if (t instanceof CameraControl.OperationCanceledException) {
-                    latchForOp1Canceled.countDown();
-                }
-            }
-        }, CameraXExecutors.mainThreadExecutor());
+        t = null;
+        try {
+            listenableFuture2.get();
+        } catch (ExecutionException e) {
+            t = e.getCause();
+        }
+        assertThat(t).isInstanceOf(CameraControl.OperationCanceledException.class);
 
-        Futures.addCallback(listenableFuture2, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-                if (t instanceof CameraControl.OperationCanceledException) {
-                    latchForOp2Canceled.countDown();
-                }
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
-        Futures.addCallback(listenableFuture3, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-                latchForOp3Succeeded.countDown();
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
-        assertTrue(latchForOp1Canceled.await(500, TimeUnit.MILLISECONDS));
-        assertTrue(latchForOp2Canceled.await(500, TimeUnit.MILLISECONDS));
-        assertTrue(latchForOp3Succeeded.await(500, TimeUnit.MILLISECONDS));
+        // Future 3 should have succeeded. Should not throw.
+        listenableFuture3.get();
     }
 
     @Test
-    public void setZoomRatioAndPercentage_mixedOperation() throws InterruptedException {
-        CountDownLatch latchForOp1Canceled = new CountDownLatch(1);
-        CountDownLatch latchForOp2Canceled = new CountDownLatch(1);
-        CountDownLatch latchForOp3Succeeded = new CountDownLatch(1);
+    public void setZoomRatioAndPercentage_mixedOperation()
+            throws InterruptedException, ExecutionException {
         ListenableFuture<Void> listenableFuture = mZoomControl.setZoomRatio(2f);
         ListenableFuture<Void> listenableFuture2 = mZoomControl.setLinearZoom(0.1f);
         ListenableFuture<Void> listenableFuture3 = mZoomControl.setZoomRatio(4f);
+        // setZoomRatio/setLinearZoom can be called from any thread and posts to executor, so idle
+        // our executor. Since multiple calls to setZoomRatio/setLinearZoom  are posted in order,
+        // we only need to idle once to run all of them.
+        shadowOf(getMainLooper()).idle();
 
         TotalCaptureResult result = mockCaptureResult(getCropRectByRatio(4.0f));
+        // Calling onCaptureResult directly from executor thread (main thread). No need to idle.
         mCaptureResultListener.onCaptureResult(result);
 
-        Futures.addCallback(listenableFuture, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-            }
+        assertThat(listenableFuture.isDone()).isTrue();
+        assertThat(listenableFuture2.isDone()).isTrue();
+        assertThat(listenableFuture3.isDone()).isTrue();
+        // Futures 1 and 2 should have failed.
+        Throwable t = null;
+        try {
+            listenableFuture.get();
+        } catch (ExecutionException e) {
+            t = e.getCause();
+        }
+        assertThat(t).isInstanceOf(CameraControl.OperationCanceledException.class);
 
-            @Override
-            public void onFailure(Throwable t) {
-                if (t instanceof CameraControl.OperationCanceledException) {
-                    latchForOp1Canceled.countDown();
-                }
-            }
-        }, CameraXExecutors.mainThreadExecutor());
+        t = null;
+        try {
+            listenableFuture2.get();
+        } catch (ExecutionException e) {
+            t = e.getCause();
+        }
+        assertThat(t).isInstanceOf(CameraControl.OperationCanceledException.class);
 
-        Futures.addCallback(listenableFuture2, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-                if (t instanceof CameraControl.OperationCanceledException) {
-                    latchForOp2Canceled.countDown();
-                }
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
-        Futures.addCallback(listenableFuture3, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-                latchForOp3Succeeded.countDown();
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
-        assertTrue(latchForOp1Canceled.await(500, TimeUnit.MILLISECONDS));
-        assertTrue(latchForOp2Canceled.await(500, TimeUnit.MILLISECONDS));
-        assertTrue(latchForOp3Succeeded.await(500, TimeUnit.MILLISECONDS));
+        // Future 3 should have succeeded. Should not throw.
+        listenableFuture3.get();
     }
 
     @Test
-    public void setZoomRatio_whenInActive_operationCanceled() {
+    public void setZoomRatio_whenInActive_operationCanceled() throws InterruptedException {
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         mZoomControl.setActive(false);
         ListenableFuture<Void> listenableFuture = mZoomControl.setZoomRatio(1.0f);
+        // setZoomRatio can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
 
+        assertThat(listenableFuture.isDone()).isTrue();
+        Throwable t = null;
         try {
-            listenableFuture.get(1000, TimeUnit.MILLISECONDS);
+            listenableFuture.get();
         } catch (ExecutionException e) {
-            if (e.getCause() instanceof CameraControl.OperationCanceledException) {
-                assertTrue(true);
-                return;
-            }
-        } catch (Exception e) {
+            t = e.getCause();
         }
-
-        fail();
+        assertThat(t).isInstanceOf(CameraControl.OperationCanceledException.class);
     }
 
     @Test
-    public void setLinearZoom_whenInActive_operationCanceled() {
+    public void setLinearZoom_whenInActive_operationCanceled() throws InterruptedException {
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         mZoomControl.setActive(false);
         ListenableFuture<Void> listenableFuture = mZoomControl.setLinearZoom(0f);
+        // setZoomRatio can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
 
+        assertThat(listenableFuture.isDone()).isTrue();
+        Throwable t = null;
         try {
-            listenableFuture.get(1000, TimeUnit.MILLISECONDS);
+            listenableFuture.get();
         } catch (ExecutionException e) {
-            if (e.getCause() instanceof CameraControl.OperationCanceledException) {
-                assertTrue(true);
-                return;
-            }
-        } catch (Exception e) {
+            t = e.getCause();
         }
-
-        fail();
+        assertThat(t).isInstanceOf(CameraControl.OperationCanceledException.class);
     }
 
     @Test
     public void setZoomRatio_afterInActive_operationCanceled() throws InterruptedException {
-        CountDownLatch latch = new CountDownLatch(1);
         ListenableFuture<Void> listenableFuture = mZoomControl.setZoomRatio(2.0f);
-        Futures.addCallback(listenableFuture, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-                if (t instanceof CameraControl.OperationCanceledException) {
-                    latch.countDown();
-                }
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
+        // setZoomRatio can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         mZoomControl.setActive(false);
 
-        assertTrue(latch.await(500, TimeUnit.MILLISECONDS));
+        assertThat(listenableFuture.isDone()).isTrue();
+        Throwable t = null;
+        try {
+            listenableFuture.get();
+        } catch (ExecutionException e) {
+            t = e.getCause();
+        }
+        assertThat(t).isInstanceOf(CameraControl.OperationCanceledException.class);
     }
 
     @Test
     public void setLinearZoom_afterInActive_operationCanceled() throws InterruptedException {
-        CountDownLatch latch = new CountDownLatch(1);
         ListenableFuture<Void> listenableFuture = mZoomControl.setLinearZoom(0.3f);
-        Futures.addCallback(listenableFuture, new FutureCallback<Void>() {
-            @Override
-            public void onSuccess(@Nullable Void result) {
-            }
-
-            @Override
-            public void onFailure(Throwable t) {
-                if (t instanceof CameraControl.OperationCanceledException) {
-                    latch.countDown();
-                }
-            }
-        }, CameraXExecutors.mainThreadExecutor());
-
+        // setZoomRatio can be called from any thread and posts to executor, so idle our executor.
+        shadowOf(getMainLooper()).idle();
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         mZoomControl.setActive(false);
 
-        assertTrue(latch.await(500, TimeUnit.MILLISECONDS));
+        assertThat(listenableFuture.isDone()).isTrue();
+        Throwable t = null;
+        try {
+            listenableFuture.get();
+        } catch (ExecutionException e) {
+            t = e.getCause();
+        }
+        assertThat(t).isInstanceOf(CameraControl.OperationCanceledException.class);
     }
 
     @Test
@@ -478,10 +399,13 @@
                 mock(CameraControlInternal.ControlUpdateCallback.class));
         ZoomControl zoomControl = new ZoomControl(mCamera2CameraControl, cameraCharacteristics,
                 CameraXExecutors.mainThreadExecutor());
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         zoomControl.setActive(true);
 
+        // LiveData is updated synchronously. No need to idle.
         zoomControl.setZoomRatio(3.0f);
-        assertThat(zoomControl.getZoomState().getValue().getZoomRatio()).isEqualTo(1.0f);
+        assertThat(Objects.requireNonNull(
+                zoomControl.getZoomState().getValue()).getZoomRatio()).isEqualTo(1.0f);
         assertThat(zoomControl.getZoomState().getValue().getLinearZoom()).isEqualTo(0.0f);
     }
 
@@ -500,10 +424,13 @@
                 mock(CameraControlInternal.ControlUpdateCallback.class));
         ZoomControl zoomControl = new ZoomControl(mCamera2CameraControl, cameraCharacteristics,
                 CameraXExecutors.mainThreadExecutor());
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         zoomControl.setActive(true);
 
+        // LiveData is updated synchronously. No need to idle.
         zoomControl.setZoomRatio(0.2f);
-        assertThat(zoomControl.getZoomState().getValue().getZoomRatio()).isEqualTo(1.0f);
+        assertThat(Objects.requireNonNull(
+                zoomControl.getZoomState().getValue()).getZoomRatio()).isEqualTo(1.0f);
         assertThat(zoomControl.getZoomState().getValue().getLinearZoom()).isEqualTo(0.0f);
     }
 
@@ -523,10 +450,13 @@
                 mock(CameraControlInternal.ControlUpdateCallback.class));
         ZoomControl zoomControl = new ZoomControl(mCamera2CameraControl, cameraCharacteristics,
                 CameraXExecutors.mainThreadExecutor());
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         zoomControl.setActive(true);
 
+        // LiveData is updated synchronously. No need to idle.
         zoomControl.setLinearZoom(0.4f);
-        assertThat(zoomControl.getZoomState().getValue().getZoomRatio()).isEqualTo(1.0f);
+        assertThat(Objects.requireNonNull(
+                zoomControl.getZoomState().getValue()).getZoomRatio()).isEqualTo(1.0f);
         // percentage is updated correctly but the zoomRatio is always 1.0f if zoom not supported.
         assertThat(zoomControl.getZoomState().getValue().getLinearZoom()).isEqualTo(0.4f);
     }
@@ -546,11 +476,14 @@
                 mock(CameraControlInternal.ControlUpdateCallback.class));
         ZoomControl zoomControl = new ZoomControl(mCamera2CameraControl, cameraCharacteristics,
                 CameraXExecutors.mainThreadExecutor());
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         zoomControl.setActive(true);
 
+        // LiveData is updated synchronously. No need to idle.
         zoomControl.setLinearZoom(0.3f);
         zoomControl.setLinearZoom(-0.2f);
-        assertThat(zoomControl.getZoomState().getValue().getZoomRatio()).isEqualTo(1.0f);
+        assertThat(Objects.requireNonNull(
+                zoomControl.getZoomState().getValue()).getZoomRatio()).isEqualTo(1.0f);
         // percentage not changed but the zoomRatio is always 1.0f if zoom not supported.
         assertThat(zoomControl.getZoomState().getValue().getLinearZoom()).isEqualTo(0.3f);
     }
@@ -570,11 +503,14 @@
                 mock(CameraControlInternal.ControlUpdateCallback.class));
         ZoomControl zoomControl = new ZoomControl(mCamera2CameraControl, cameraCharacteristics,
                 CameraXExecutors.mainThreadExecutor());
+        // setActive() is called from executor thread (main thread in this case). No need to idle.
         zoomControl.setActive(true);
 
+        // LiveData is updated synchronously. No need to idle.
         zoomControl.setLinearZoom(0.3f);
         zoomControl.setLinearZoom(1.2f);
-        assertThat(zoomControl.getZoomState().getValue().getZoomRatio()).isEqualTo(1.0f);
+        assertThat(Objects.requireNonNull(
+                zoomControl.getZoomState().getValue()).getZoomRatio()).isEqualTo(1.0f);
         // percentage not changed but the zoomRatio is always 1.0f if zoom not supported.
         assertThat(zoomControl.getZoomState().getValue().getLinearZoom()).isEqualTo(0.3f);
     }
diff --git a/camera/camera-core/src/test/java/androidx/camera/core/ImageAnalysisNonBlockingAnalyzerTest.java b/camera/camera-core/src/test/java/androidx/camera/core/ImageAnalysisNonBlockingAnalyzerTest.java
index 42ba496..ecd8e5e 100644
--- a/camera/camera-core/src/test/java/androidx/camera/core/ImageAnalysisNonBlockingAnalyzerTest.java
+++ b/camera/camera-core/src/test/java/androidx/camera/core/ImageAnalysisNonBlockingAnalyzerTest.java
@@ -16,6 +16,8 @@
 
 package androidx.camera.core;
 
+import static android.os.Looper.getMainLooper;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertEquals;
@@ -26,6 +28,7 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.os.Build;
 
@@ -112,6 +115,8 @@
 
         mImageAnalysisNonBlockingAnalyzer.onImageAvailable(mImageReaderProxy);
 
+        shadowOf(getMainLooper()).idle();
+
         ArgumentCaptor<ImageProxy> imageProxyArgumentCaptor =
                 ArgumentCaptor.forClass(ImageProxy.class);
         verify(mAnalyzer, times(1)).analyze(
diff --git a/camera/camera-core/src/test/java/androidx/camera/core/ImageAnalysisTest.java b/camera/camera-core/src/test/java/androidx/camera/core/ImageAnalysisTest.java
index 1ca30c1..90fc0b2 100644
--- a/camera/camera-core/src/test/java/androidx/camera/core/ImageAnalysisTest.java
+++ b/camera/camera-core/src/test/java/androidx/camera/core/ImageAnalysisTest.java
@@ -18,6 +18,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.robolectric.Shadows.shadowOf;
+
 import android.content.Context;
 import android.graphics.Rect;
 import android.os.Build;
@@ -48,8 +50,6 @@
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.annotation.Config;
 import org.robolectric.annotation.internal.DoNotInstrument;
-import org.robolectric.shadow.api.Shadow;
-import org.robolectric.shadows.ShadowLooper;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -88,10 +88,14 @@
     public void setUp() throws ExecutionException, InterruptedException {
         mCallbackThread = new HandlerThread("Callback");
         mCallbackThread.start();
+        // Explicitly pause callback thread since we will control execution manually in tests
+        shadowOf(mCallbackThread.getLooper()).pause();
         mCallbackHandler = new Handler(mCallbackThread.getLooper());
 
         mBackgroundThread = new HandlerThread("Background");
         mBackgroundThread.start();
+        // Explicitly pause background thread since we will control execution manually in tests
+        shadowOf(mBackgroundThread.getLooper()).pause();
         mBackgroundHandler = new Handler(mBackgroundThread.getLooper());
         mBackgroundExecutor = CameraXExecutors.newHandlerExecutor(mBackgroundHandler);
 
@@ -327,6 +331,6 @@
      * @param handler the {@link Handler} to flush.
      */
     private static void flushHandler(Handler handler) {
-        ((ShadowLooper) Shadow.extract(handler.getLooper())).idle();
+        shadowOf(handler.getLooper()).idle();
     }
 }
diff --git a/camera/camera-core/src/test/java/androidx/camera/core/ImageCaptureTest.kt b/camera/camera-core/src/test/java/androidx/camera/core/ImageCaptureTest.kt
index 960f00f..f559c0c1 100644
--- a/camera/camera-core/src/test/java/androidx/camera/core/ImageCaptureTest.kt
+++ b/camera/camera-core/src/test/java/androidx/camera/core/ImageCaptureTest.kt
@@ -22,6 +22,7 @@
 import android.os.Build
 import android.os.Handler
 import android.os.HandlerThread
+import android.os.Looper.getMainLooper
 import android.util.Pair
 import android.util.Rational
 import android.view.Surface
@@ -46,7 +47,6 @@
 import androidx.concurrent.futures.ResolvableFuture
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.filters.MediumTest
-import com.google.common.truth.Truth
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
@@ -60,6 +60,7 @@
 import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
 import org.robolectric.RobolectricTestRunner
+import org.robolectric.Shadows.shadowOf
 import org.robolectric.annotation.Config
 import org.robolectric.annotation.internal.DoNotInstrument
 import org.robolectric.shadow.api.Shadow
@@ -115,6 +116,8 @@
         CameraX.initialize(context, cameraXConfig).get()
         callbackThread = HandlerThread("Callback")
         callbackThread.start()
+        // Explicitly pause callback thread since we will control execution manually in tests
+        shadowOf(callbackThread.looper).pause()
         callbackHandler = Handler(callbackThread.looper)
         executor = CameraXExecutors.newHandlerExecutor(callbackHandler)
     }
@@ -138,6 +141,7 @@
         imageCapture.takePicture(executor, onImageCapturedCallback)
         // Send fake image.
         fakeImageReaderProxy?.triggerImageAvailable(TagBundle.create(Pair("TagBundleKey", 0)), 0)
+        shadowOf(getMainLooper()).idle()
         flushHandler(callbackHandler)
 
         // Assert.
@@ -182,11 +186,12 @@
         imageCapture.takePicture(executor, onImageCapturedCallback)
         // Send fake image.
         fakeImageReaderProxy?.triggerImageAvailable(TagBundle.create(Pair("TagBundleKey", 0)), 0)
+        shadowOf(getMainLooper()).idle()
         flushHandler(callbackHandler)
 
         // Assert.
-        Truth.assertThat(capturedImage!!.width).isEqualTo(fakeImageReaderProxy?.width)
-        Truth.assertThat(capturedImage!!.height).isEqualTo(fakeImageReaderProxy?.height)
+        assertThat(capturedImage!!.width).isEqualTo(fakeImageReaderProxy?.width)
+        assertThat(capturedImage!!.height).isEqualTo(fakeImageReaderProxy?.height)
     }
 
     @Test
diff --git a/camera/camera-core/src/test/java/androidx/camera/core/PreviewTest.kt b/camera/camera-core/src/test/java/androidx/camera/core/PreviewTest.kt
index 273f27b..b478e33 100644
--- a/camera/camera-core/src/test/java/androidx/camera/core/PreviewTest.kt
+++ b/camera/camera-core/src/test/java/androidx/camera/core/PreviewTest.kt
@@ -19,6 +19,7 @@
 import android.content.Context
 import android.graphics.Rect
 import android.os.Build
+import android.os.Looper.getMainLooper
 import android.util.Rational
 import android.util.Size
 import android.view.Surface
@@ -40,6 +41,7 @@
 import org.robolectric.RobolectricTestRunner
 import org.robolectric.annotation.Config
 import org.robolectric.annotation.internal.DoNotInstrument
+import org.robolectric.Shadows.shadowOf
 import java.util.Collections
 import java.util.concurrent.ExecutionException
 
@@ -126,18 +128,20 @@
                 .getApplicationContext<Context>(), CameraSelector.DEFAULT_BACK_CAMERA
         )
         cameraUseCaseAdapter.addUseCases(Collections.singleton<UseCase>(preview))
-        InstrumentationRegistry.getInstrumentation().waitForIdleSync()
+
         // Unsent pending SurfaceRequest created by pipeline.
         val pendingSurfaceRequest = preview.mPendingSurfaceRequest
         var receivedSurfaceRequest: SurfaceRequest? = null
 
         // Act: set a SurfaceProvider after attachment.
         preview.setSurfaceProvider { receivedSurfaceRequest = it }
+        shadowOf(getMainLooper()).idle()
         // Assert: received a SurfaceRequest.
         assertThat(receivedSurfaceRequest).isSameInstanceAs(pendingSurfaceRequest)
 
         // Act: set a different SurfaceProvider.
         preview.setSurfaceProvider { receivedSurfaceRequest = it }
+        shadowOf(getMainLooper()).idle()
         // Assert: received a different SurfaceRequest.
         assertThat(receivedSurfaceRequest).isNotSameInstanceAs(pendingSurfaceRequest)
     }
diff --git a/camera/camera-core/src/test/java/androidx/camera/core/ProcessingImageReaderTest.java b/camera/camera-core/src/test/java/androidx/camera/core/ProcessingImageReaderTest.java
index d27dccf..c40e6d8 100644
--- a/camera/camera-core/src/test/java/androidx/camera/core/ProcessingImageReaderTest.java
+++ b/camera/camera-core/src/test/java/androidx/camera/core/ProcessingImageReaderTest.java
@@ -16,11 +16,14 @@
 
 package androidx.camera.core;
 
+import static android.os.Looper.getMainLooper;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.robolectric.Shadows.shadowOf;
 
 import android.graphics.ImageFormat;
 import android.os.AsyncTask;
@@ -41,15 +44,18 @@
 import androidx.camera.testing.fakes.FakeImageReaderProxy;
 import androidx.test.filters.SmallTest;
 
+import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Before;
+import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.android.util.concurrent.PausedExecutorService;
 import org.robolectric.annotation.Config;
 import org.robolectric.annotation.internal.DoNotInstrument;
-import org.robolectric.shadows.ShadowLooper;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -58,6 +64,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
+@SuppressWarnings("UnstableApiUsage") // Needed because PausedExecutorService is marked @Beta
 @SmallTest
 @RunWith(RobolectricTestRunner.class)
 @DoNotInstrument
@@ -87,7 +94,7 @@
 
         }
     };
-
+    private static PausedExecutorService sPausedExecutor;
     private final CaptureStage mCaptureStage0 = new FakeCaptureStage(CAPTURE_ID_0, null);
     private final CaptureStage mCaptureStage1 = new FakeCaptureStage(CAPTURE_ID_1, null);
     private final CaptureStage mCaptureStage2 = new FakeCaptureStage(CAPTURE_ID_2, null);
@@ -97,6 +104,16 @@
     private CaptureBundle mCaptureBundle;
     private String mTagBundleKey;
 
+    @BeforeClass
+    public static void setUpClass() {
+        sPausedExecutor = new PausedExecutorService();
+    }
+
+    @AfterClass
+    public static void tearDownClass() {
+        sPausedExecutor.shutdown();
+    }
+
     @Before
     public void setUp() {
         mCaptureBundle = CaptureBundles.createCaptureBundle(mCaptureStage0, mCaptureStage1);
@@ -104,13 +121,19 @@
         mMetadataImageReader = new MetadataImageReader(mImageReaderProxy);
     }
 
+    @After
+    public void cleanUp() {
+        // Ensure the PausedExecutorService is drained
+        sPausedExecutor.runAll();
+    }
+
     @Test
     public void canSetFuturesInSettableImageProxyBundle()
             throws InterruptedException, TimeoutException, ExecutionException {
         // Sets the callback from ProcessingImageReader to start processing
         CaptureProcessor captureProcessor = mock(CaptureProcessor.class);
         ProcessingImageReader processingImageReader = new ProcessingImageReader(
-                mMetadataImageReader, AsyncTask.THREAD_POOL_EXECUTOR, mCaptureBundle,
+                mMetadataImageReader, sPausedExecutor, mCaptureBundle,
                 captureProcessor);
         processingImageReader.setOnImageAvailableListener(mock(
                 ImageReaderProxy.OnImageAvailableListener.class),
@@ -145,12 +168,15 @@
             triggerImageAvailable(id, captureIdToTime.get(id));
         }
 
-        // Ensure all posted tasks finish running
-        ShadowLooper.runUiThreadTasks();
+        // Ensure tasks are posted to the processing executor
+        shadowOf(getMainLooper()).idle();
+
+        // Run processing
+        sPausedExecutor.runAll();
 
         ArgumentCaptor<ImageProxyBundle> imageProxyBundleCaptor =
                 ArgumentCaptor.forClass(ImageProxyBundle.class);
-        verify(captureProcessor, timeout(3000).times(1)).process(imageProxyBundleCaptor.capture());
+        verify(captureProcessor, times(1)).process(imageProxyBundleCaptor.capture());
         assertThat(imageProxyBundleCaptor.getValue()).isNotNull();
 
         // CaptureProcessor.process should be called once all ImageProxies on the
@@ -187,6 +213,9 @@
             triggerImageAvailable(idTimestamp.getKey(), idTimestamp.getValue());
         }
 
+        // Ensure tasks are posted to the processing executor
+        shadowOf(getMainLooper()).idle();
+
         // Wait for CaptureProcessor.process() to start so that it is in the middle of processing
         assertThat(waitingCaptureProcessor.waitForProcessingToStart(3000)).isTrue();
 
@@ -194,7 +223,7 @@
 
         // Allow the CaptureProcessor to continue processing. Calling finishProcessing() will
         // cause the CaptureProcessor to start accessing the ImageProxy. If the ImageProxy has
-        // already been closed then
+        // already been closed then we will time out at waitForProcessingToComplete().
         waitingCaptureProcessor.finishProcessing();
 
         // The processing will only complete if no exception was thrown during the processing
@@ -202,11 +231,13 @@
         assertThat(waitingCaptureProcessor.waitForProcessingToComplete(3000)).isTrue();
     }
 
+    // Tests that a ProcessingImageReader can be closed while in the process of receiving
+    // ImageProxies for an ImageProxyBundle.
     @Test
     public void closeImageHalfway() throws InterruptedException {
         // Sets the callback from ProcessingImageReader to start processing
         ProcessingImageReader processingImageReader = new ProcessingImageReader(
-                mMetadataImageReader, AsyncTask.THREAD_POOL_EXECUTOR, mCaptureBundle,
+                mMetadataImageReader, sPausedExecutor, mCaptureBundle,
                 NOOP_PROCESSOR);
         processingImageReader.setOnImageAvailableListener(mock(
                 ImageReaderProxy.OnImageAvailableListener.class),
@@ -216,9 +247,12 @@
         mTagBundleKey = processingImageReader.getTagBundleKey();
         triggerImageAvailable(CAPTURE_ID_0, TIMESTAMP_0);
 
-        processingImageReader.close();
+        // Ensure the first image is received by the ProcessingImageReader
+        shadowOf(getMainLooper()).idle();
 
-        triggerImageAvailable(CAPTURE_ID_1, TIMESTAMP_1);
+        // The ProcessingImageReader is closed after receiving the first image, but before
+        // receiving enough images for the entire ImageProxyBundle.
+        processingImageReader.close();
 
         assertThat(mImageReaderProxy.isClosed()).isTrue();
     }
@@ -263,15 +297,15 @@
     private static class WaitingCaptureProcessor implements CaptureProcessor {
         // Block processing so that the ProcessingImageReader can be closed before the
         // CaptureProcessor has finished accessing the ImageProxy and ImageProxyBundle
-        private CountDownLatch mProcessingLatch = new CountDownLatch(1);
+        private final CountDownLatch mProcessingLatch = new CountDownLatch(1);
 
         // To wait for processing to start. This makes sure that the ProcessingImageReader can be
         // closed after processing has started
-        private CountDownLatch mProcessingStartLatch = new CountDownLatch(1);
+        private final CountDownLatch mProcessingStartLatch = new CountDownLatch(1);
 
         // Block processing from completing. This ensures that the CaptureProcessor has finished
         // accessing the ImageProxy and ImageProxyBundle successfully.
-        private CountDownLatch mProcessingComplete = new CountDownLatch(1);
+        private final CountDownLatch mProcessingComplete = new CountDownLatch(1);
 
         WaitingCaptureProcessor() {
         }
diff --git a/camera/camera-core/src/test/java/androidx/camera/core/impl/CameraRepositoryTest.java b/camera/camera-core/src/test/java/androidx/camera/core/impl/CameraRepositoryTest.java
index 701319c..7a6b809 100644
--- a/camera/camera-core/src/test/java/androidx/camera/core/impl/CameraRepositoryTest.java
+++ b/camera/camera-core/src/test/java/androidx/camera/core/impl/CameraRepositoryTest.java
@@ -16,8 +16,12 @@
 
 package androidx.camera.core.impl;
 
+import static android.os.Looper.getMainLooper;
+
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.robolectric.Shadows.shadowOf;
+
 import android.os.Build;
 
 import androidx.camera.core.CameraSelector;
@@ -35,7 +39,6 @@
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.annotation.Config;
 import org.robolectric.annotation.internal.DoNotInstrument;
-import org.robolectric.shadows.ShadowLooper;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -118,12 +121,15 @@
         ListenableFuture<Void> deinitFuture = mCameraRepository.deinit();
 
         // Needed since FakeCamera uses LiveDataObservable
-        ShadowLooper.runUiThreadTasks();
+        shadowOf(getMainLooper()).idle();
 
         assertThat(deinitFuture.isDone()).isTrue();
         for (CameraInternal cameraInternal : cameraInternals) {
-            assertThat(cameraInternal.getCameraState().fetchData().get()).isEqualTo(
-                    CameraInternal.State.RELEASED);
+            ListenableFuture<CameraInternal.State> stateFuture =
+                    cameraInternal.getCameraState().fetchData();
+            // Needed since FakeCamera uses LiveDataObservable
+            shadowOf(getMainLooper()).idle();
+            assertThat(stateFuture.get()).isEqualTo(CameraInternal.State.RELEASED);
         }
     }
 }
diff --git a/camera/camera-core/src/test/java/androidx/camera/core/impl/utils/ExifTest.java b/camera/camera-core/src/test/java/androidx/camera/core/impl/utils/ExifTest.java
index 7a65d76..8fa4dc5 100644
--- a/camera/camera-core/src/test/java/androidx/camera/core/impl/utils/ExifTest.java
+++ b/camera/camera-core/src/test/java/androidx/camera/core/impl/utils/ExifTest.java
@@ -20,6 +20,7 @@
 
 import android.location.Location;
 import android.os.Build;
+import android.os.SystemClock;
 
 import androidx.test.filters.SmallTest;
 
@@ -34,6 +35,7 @@
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.time.Duration;
 
 @SmallTest
 @RunWith(RobolectricTestRunner.class)
@@ -156,13 +158,12 @@
 
     @Test
     public void attachedTimestampUsesSystemWallTime() {
-        long beforeTimestamp = System.currentTimeMillis();
+        long beforeTimestamp = SystemClock.uptimeMillis();
+        ShadowSystemClock.advanceBy(Duration.ofMillis(100));
 
-        // The Exif class is instrumented since it's in the androidx.* namespace.
-        // Set the ShadowSystemClock to match the real system clock.
-        ShadowSystemClock.setNanoTime(System.currentTimeMillis() * 1000 * 1000);
         mExif.attachTimestamp();
-        long afterTimestamp = System.currentTimeMillis();
+        ShadowSystemClock.advanceBy(Duration.ofMillis(100));
+        long afterTimestamp = SystemClock.uptimeMillis();
 
         // Check that the attached timestamp is in the closed range [beforeTimestamp,
         // afterTimestamp].
diff --git a/camera/camera-core/src/test/java/androidx/camera/core/impl/utils/MainThreadAsyncHandlerTest.java b/camera/camera-core/src/test/java/androidx/camera/core/impl/utils/MainThreadAsyncHandlerTest.java
index 3917a8d..f8e2ad4 100644
--- a/camera/camera-core/src/test/java/androidx/camera/core/impl/utils/MainThreadAsyncHandlerTest.java
+++ b/camera/camera-core/src/test/java/androidx/camera/core/impl/utils/MainThreadAsyncHandlerTest.java
@@ -24,8 +24,6 @@
 
 import androidx.test.filters.SmallTest;
 
-import org.junit.After;
-import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.robolectric.RobolectricTestRunner;
@@ -41,16 +39,6 @@
 @Config(minSdk = Build.VERSION_CODES.LOLLIPOP)
 public class MainThreadAsyncHandlerTest {
 
-    @Before
-    public void setUp() {
-        ShadowLooper.pauseMainLooper();
-    }
-
-    @After
-    public void tearDown() {
-        ShadowLooper.idleMainLooperConstantly(true);
-    }
-
     @Test
     public void canPostTaskToMainLooper() {
         Handler handler = MainThreadAsyncHandler.getInstance();
diff --git a/camera/camera-extensions/build.gradle b/camera/camera-extensions/build.gradle
index 0633ff5..c6415d4 100644
--- a/camera/camera-extensions/build.gradle
+++ b/camera/camera-extensions/build.gradle
@@ -39,7 +39,7 @@
     testImplementation(JUNIT)
     testImplementation(MOCKITO_CORE)
     testImplementation(ROBOLECTRIC)
-
+    testImplementation(TRUTH)
     testImplementation project(":camera:camera-testing")
     testImplementation(project(":camera:camera-extensions-stub"))
     // To use the extensions-stub for testing directly.
diff --git a/camera/camera-extensions/src/test/java/androidx/camera/extensions/VersionTest.java b/camera/camera-extensions/src/test/java/androidx/camera/extensions/VersionTest.java
index fb6859e..31df428 100644
--- a/camera/camera-extensions/src/test/java/androidx/camera/extensions/VersionTest.java
+++ b/camera/camera-extensions/src/test/java/androidx/camera/extensions/VersionTest.java
@@ -16,15 +16,7 @@
 
 package androidx.camera.extensions;
 
-import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThan;
-import static org.hamcrest.Matchers.lessThan;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertThat;
-import static org.junit.Assert.assertTrue;
+import static com.google.common.truth.Truth.assertThat;
 
 import org.junit.Test;
 
@@ -39,60 +31,60 @@
 
         Version version2 = Version.create(2, 0, 0, "test");
 
-        assertTrue(version1.equals(version1_description));
-        assertFalse(version1.equals(version1_patch));
-        assertFalse(version1.equals(version1_minor));
-        assertFalse(version1.equals(version2));
+        assertThat(version1.equals(version1_description)).isTrue();
+        assertThat(version1.equals(version1_patch)).isFalse();
+        assertThat(version1.equals(version1_minor)).isFalse();
+        assertThat(version1.equals(version2)).isFalse();
 
-        assertThat(version1.compareTo(version1_patch), lessThan(0));
-        assertThat(version1.compareTo(version1_description), equalTo(0));
-        assertThat(version1.compareTo(version1_minor), lessThan(0));
-        assertThat(version1.compareTo(version2), lessThan(0));
+        assertThat(version1.compareTo(version1_patch)).isLessThan(0);
+        assertThat(version1.compareTo(version1_description)).isEqualTo(0);
+        assertThat(version1.compareTo(version1_minor)).isLessThan(0);
+        assertThat(version1.compareTo(version2)).isLessThan(0);
 
-        assertThat(version1.compareTo(1), equalTo(0));
-        assertThat(version1.compareTo(2), lessThan(0));
-        assertThat(version1.compareTo(0), greaterThan(0));
+        assertThat(version1.compareTo(1)).isEqualTo(0);
+        assertThat(version1.compareTo(2)).isLessThan(0);
+        assertThat(version1.compareTo(0)).isGreaterThan(0);
 
-        assertThat(version1.compareTo(1, 0), equalTo(0));
-        assertThat(version1.compareTo(1, 1), lessThan(0));
-        assertThat(version1_minor.compareTo(1, 0), greaterThan(0));
+        assertThat(version1.compareTo(1, 0)).isEqualTo(0);
+        assertThat(version1.compareTo(1, 1)).isLessThan(0);
+        assertThat(version1_minor.compareTo(1, 0)).isGreaterThan(0);
 
-        assertThat(version1.compareTo(2, 0), lessThan(0));
+        assertThat(version1.compareTo(2, 0)).isLessThan(0);
     }
 
     @Test
     public void testParseStringVersion() {
 
         Version version1 = Version.parse("1.2.3-description");
-        assertNotNull(version1);
-        assertEquals(version1.getMajor(), 1);
-        assertEquals(version1.getMinor(), 2);
-        assertEquals(version1.getPatch(), 3);
-        assertEquals(version1.getDescription(), "description");
+        assertThat(version1).isNotNull();
+        assertThat(version1.getMajor()).isEqualTo(1);
+        assertThat(version1.getMinor()).isEqualTo(2);
+        assertThat(version1.getPatch()).isEqualTo(3);
+        assertThat(version1.getDescription()).isEqualTo("description");
 
         Version version2 = Version.parse("4.5.6");
-        assertNotNull(version2);
-        assertEquals(version2.getDescription(), "");
+        assertThat(version2).isNotNull();
+        assertThat(version2.getDescription()).isEqualTo("");
 
         Version version3 = Version.parse("01.002.0003");
-        assertNotNull(version3);
-        assertEquals(version3.getMajor(), 1);
-        assertEquals(version3.getMinor(), 2);
-        assertEquals(version3.getPatch(), 3);
+        assertThat(version3).isNotNull();
+        assertThat(version3.getMajor()).isEqualTo(1);
+        assertThat(version3.getMinor()).isEqualTo(2);
+        assertThat(version3.getPatch()).isEqualTo(3);
 
 
         // Test invalid input version string.
-        assertNull(Version.parse("1.0"));
-        assertNull(Version.parse("1. 0.0"));
-        assertNull(Version.parse("1..0"));
-        assertNull(Version.parse("1.0.a"));
-        assertNull(Version.parse("1.0.0."));
-        assertNull(Version.parse("1.0.0.description"));
+        assertThat(Version.parse("1.0")).isNull();
+        assertThat(Version.parse("1. 0.0")).isNull();
+        assertThat(Version.parse("1..0")).isNull();
+        assertThat(Version.parse("1.0.a")).isNull();
+        assertThat(Version.parse("1.0.0.")).isNull();
+        assertThat(Version.parse("1.0.0.description")).isNull();
 
-        assertNull(Version.parse("1.0.0.0"));
-        assertNull(Version.parse("1.0.-0"));
-        assertNull(Version.parse("1.0.-0"));
-        assertNull(Version.parse("(1.0.0)"));
-        assertNull(Version.parse(" 1.0.0 "));
+        assertThat(Version.parse("1.0.0.0")).isNull();
+        assertThat(Version.parse("1.0.-0")).isNull();
+        assertThat(Version.parse("1.0.-0")).isNull();
+        assertThat(Version.parse("(1.0.0)")).isNull();
+        assertThat(Version.parse(" 1.0.0 ")).isNull();
     }
 }