Fix flaky tests

Author: tejasd

firebase-perf/src/test/java/com/google/firebase/perf/session/gauges/GaugeManagerTest.java

@@ -45,6 +45,7 @@
 import com.google.testing.timing.FakeScheduledExecutorService;
 import java.util.Random;
 import java.util.concurrent.TimeUnit;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -126,6 +127,11 @@ public void setUp() {
             new Lazy<>(() -> fakeMemoryGaugeCollector));
   }
 
+  @After
+  public void tearDown() {
+    shadowOf(Looper.getMainLooper()).idle();
+  }
+
   @Test
   public void testStartCollectingGaugesStartsCollectingMetricsDefault() {
     PerfSession fakeSession = createTestSession(1);
@@ -347,7 +353,6 @@ public void testGaugeCounterStartsAJobToConsumeTheGeneratedMetrics() {
         .isEqualTo(TIME_TO_WAIT_BEFORE_FLUSHING_GAUGES_QUEUE_MS);
 
     fakeScheduledExecutorService.simulateSleepExecutingAtMostOneTask();
-    shadowOf(Looper.getMainLooper()).idle();
 
     // Generate additional metrics, but doesn't start logging them as it hasn't met the threshold.
     generateMetricsAndIncrementCounter(5);
@@ -418,59 +423,59 @@ public void testUpdateAppStateHandlesMultipleAppStates() {
     assertThat(recordedGaugeMetric.getSessionId()).isEqualTo(testSessionId(1));
   }
 
-  //  @Test
-  //  public void testGaugeManagerHandlesMultipleSessionIds() {
-  //    PerfSession fakeSession = createTestSession(1);
-  //    fakeSession.setGaugeAndEventCollectionEnabled(true);
-  //    testGaugeManager.setApplicationProcessState(ApplicationProcessState.BACKGROUND);
-  //    testGaugeManager.startCollectingGauges(fakeSession);
-  //    GaugeCounter.INSTANCE.setGaugeManager(testGaugeManager);
-  //
-  //    // Generate metrics that don't exceed the GaugeCounter.MAX_COUNT.
-  //    generateMetricsAndIncrementCounter(10);
-  //
-  //    PerfSession updatedPerfSession = createTestSession(2);
-  //    updatedPerfSession.setGaugeAndEventCollectionEnabled(true);
-  //
-  //    // A new session and updated app state.
-  //    testGaugeManager.startCollectingGauges(updatedPerfSession);
-  //    testGaugeManager.setApplicationProcessState(ApplicationProcessState.FOREGROUND);
-  //
-  //    assertThat(fakeScheduledExecutorService.isEmpty()).isFalse();
-  //    assertThat(fakeScheduledExecutorService.getDelayToNextTask(TimeUnit.MILLISECONDS))
-  //        .isEqualTo(TIME_TO_WAIT_BEFORE_FLUSHING_GAUGES_QUEUE_MS);
-  //
-  //    fakeScheduledExecutorService.simulateSleepExecutingAtMostOneTask();
-  //    shadowOf(Looper.getMainLooper()).idle();
-  //
-  //    // Generate metrics for the new session.
-  //    generateMetricsAndIncrementCounter(26);
-  //
-  //    GaugeMetric recordedGaugeMetric =
-  //        getLastRecordedGaugeMetric(ApplicationProcessState.BACKGROUND);
-  //
-  //    // It flushes all metrics in the ConcurrentLinkedQueues.
-  //    int recordedGaugeMetricsCount =
-  //        recordedGaugeMetric.getAndroidMemoryReadingsCount()
-  //            + recordedGaugeMetric.getCpuMetricReadingsCount();
-  //    assertThat(recordedGaugeMetricsCount).isEqualTo(10);
-  //
-  //    assertThat(recordedGaugeMetric.getSessionId()).isEqualTo(testSessionId(1));
-  //
-  //    // Simulate gauges collected in the new app state.
-  //    fakeScheduledExecutorService.simulateSleepExecutingAtMostOneTask();
-  //    shadowOf(Looper.getMainLooper()).idle();
-  //
-  //    recordedGaugeMetric = getLastRecordedGaugeMetric(ApplicationProcessState.BACKGROUND);
-  //
-  //    // Verify the metrics in the new app state.
-  //    recordedGaugeMetricsCount =
-  //        recordedGaugeMetric.getAndroidMemoryReadingsCount()
-  //            + recordedGaugeMetric.getCpuMetricReadingsCount();
-  //    assertThat(recordedGaugeMetricsCount).isEqualTo(26);
-  //
-  //    assertThat(recordedGaugeMetric.getSessionId()).isEqualTo(testSessionId(2));
-  //  }
+  @Test
+  public void testGaugeManagerHandlesMultipleSessionIds() {
+    PerfSession fakeSession = createTestSession(1);
+    fakeSession.setGaugeAndEventCollectionEnabled(true);
+    testGaugeManager.setApplicationProcessState(ApplicationProcessState.BACKGROUND);
+    testGaugeManager.startCollectingGauges(fakeSession);
+    GaugeCounter.INSTANCE.setGaugeManager(testGaugeManager);
+
+    // Generate metrics that don't exceed the GaugeCounter.MAX_COUNT.
+    generateMetricsAndIncrementCounter(10);
+
+    PerfSession updatedPerfSession = createTestSession(2);
+    updatedPerfSession.setGaugeAndEventCollectionEnabled(true);
+
+    // A new session and updated app state.
+    testGaugeManager.startCollectingGauges(updatedPerfSession);
+    testGaugeManager.setApplicationProcessState(ApplicationProcessState.FOREGROUND);
+
+    assertThat(fakeScheduledExecutorService.isEmpty()).isFalse();
+    assertThat(fakeScheduledExecutorService.getDelayToNextTask(TimeUnit.MILLISECONDS))
+        .isEqualTo(TIME_TO_WAIT_BEFORE_FLUSHING_GAUGES_QUEUE_MS);
+
+    fakeScheduledExecutorService.simulateSleepExecutingAtMostOneTask();
+    shadowOf(Looper.getMainLooper()).idle();
+
+    // Generate metrics for the new session.
+    generateMetricsAndIncrementCounter(26);
+
+    GaugeMetric recordedGaugeMetric =
+        getLastRecordedGaugeMetric(ApplicationProcessState.BACKGROUND);
+
+    // It flushes all metrics in the ConcurrentLinkedQueues.
+    int recordedGaugeMetricsCount =
+        recordedGaugeMetric.getAndroidMemoryReadingsCount()
+            + recordedGaugeMetric.getCpuMetricReadingsCount();
+    assertThat(recordedGaugeMetricsCount).isEqualTo(10);
+
+    assertThat(recordedGaugeMetric.getSessionId()).isEqualTo(testSessionId(1));
+
+    // Simulate gauges collected in the new app state.
+    fakeScheduledExecutorService.simulateSleepExecutingAtMostOneTask();
+    shadowOf(Looper.getMainLooper()).idle();
+
+    recordedGaugeMetric = getLastRecordedGaugeMetric(ApplicationProcessState.FOREGROUND);
+
+    // Verify the metrics in the new app state.
+    recordedGaugeMetricsCount =
+        recordedGaugeMetric.getAndroidMemoryReadingsCount()
+            + recordedGaugeMetric.getCpuMetricReadingsCount();
+    assertThat(recordedGaugeMetricsCount).isEqualTo(26);
+
+    assertThat(recordedGaugeMetric.getSessionId()).isEqualTo(testSessionId(2));
+  }
 
   @Test
   public void testStopCollectingGaugesStopsCollectingAllGaugeMetrics() {