Fixed tests timing out when CUJ is set to stay in app until next CUJ.

This was caused by idle time in CUJ set to hit right up to the timeout
deadline. This especially became a problem after LongevityClassRunner
put @AfterClass into @After, which resulted in @AfterClass not having
time to execute before hitting the timeout.

Bug: 144958271
Test: atest LongevityPlatformLibTests:ScheduledScenarioRunnerTest (issue
reproduced in test and verified fixed)
Change-Id: Iee3891781307953ec9f2f65ea5ece57e5c5e24a6
(cherry picked from commit 5fcd9c0bf986ae119549ec43010c081c9c5e9765)
diff --git a/libraries/health/runners/longevity/platform/src/android/platform/test/longevity/ProfileSuite.java b/libraries/health/runners/longevity/platform/src/android/platform/test/longevity/ProfileSuite.java
index 3e6bee1..2baa274 100644
--- a/libraries/health/runners/longevity/platform/src/android/platform/test/longevity/ProfileSuite.java
+++ b/libraries/health/runners/longevity/platform/src/android/platform/test/longevity/ProfileSuite.java
@@ -38,10 +38,6 @@
  * profile.
  */
 public class ProfileSuite extends LongevitySuite {
-    // An arbiturary leeway that hopefully allows the @After and @AfterClass methods of a scenario
-    // to finish execution.
-    @VisibleForTesting static final long ENDTIME_LEEWAY_MS = 3000;
-
     private static final String LOG_TAG = ProfileSuite.class.getSimpleName();
 
     // Profile instance for scheduling tests.
diff --git a/libraries/health/runners/longevity/platform/src/android/platform/test/longevity/ScheduledScenarioRunner.java b/libraries/health/runners/longevity/platform/src/android/platform/test/longevity/ScheduledScenarioRunner.java
index a263e0a..1c4043e 100644
--- a/libraries/health/runners/longevity/platform/src/android/platform/test/longevity/ScheduledScenarioRunner.java
+++ b/libraries/health/runners/longevity/platform/src/android/platform/test/longevity/ScheduledScenarioRunner.java
@@ -40,10 +40,18 @@
  * performs an idle before teardown (staying inside the app for Android CUJs).
  */
 public class ScheduledScenarioRunner extends LongevityClassRunner {
-    @VisibleForTesting static final long ENDTIME_LEEWAY_MS = 3000;
+    // A leeway to ensure that the teardown steps in @After and @AfterClass has time to finish.
+    // Regardless of test passing (in which case teardown is considered in test timeout) or failing
+    // (in which case teardown happens outside the scope of the timeout).
+    // Please note that in most cases (when the CUJ does not time out) the actual cushion for
+    // teardown is double the value below, as a cushion needs to be created inside the timeout
+    // rule and also outside of it.
+    @VisibleForTesting static final long TEARDOWN_LEEWAY_MS = 2000;
 
     private final Scenario mScenario;
     private final long mTotalTimeoutMs;
+    // Timeout after the teardown leeway is taken into account.
+    private final long mEnforcedTimeoutMs;
     private final boolean mShouldIdle;
     private final Bundle mArguments;
 
@@ -63,6 +71,9 @@
         mScenario = scenario;
         // Ensure that the timeout is non-negative.
         mTotalTimeoutMs = max(timeout, 0);
+        // Ensure that the enforced timeout is non-negative. This cushion is built in so that the
+        // CUJ still has time for teardown steps when the test portion times out.
+        mEnforcedTimeoutMs = max(mTotalTimeoutMs - TEARDOWN_LEEWAY_MS, 0);
         mShouldIdle = shouldIdle;
         mArguments = arguments;
     }
@@ -71,7 +82,7 @@
     protected List<TestRule> getTestRules(Object target) {
         List<TestRule> rules = super.getTestRules(target);
         // Ensure that the timeout rule has a non-negative timeout.
-        rules.add(0, Timeout.millis(max(mTotalTimeoutMs - ENDTIME_LEEWAY_MS, 0)));
+        rules.add(0, Timeout.millis(mEnforcedTimeoutMs));
         return rules;
     }
 
@@ -85,9 +96,9 @@
                             // Run the underlying test and report exceptions.
                             statement.evaluate();
                         } finally {
-                            // If there is time left for idling (i.e. more than ENDTIME_LEEWAY_MS),
+                            // If there is time left for idling (i.e. more than TEARDOWN_LEEWAY_MS),
                             // and the scenario is set to stay in app, idle for the remainder of
-                            // its timeout window until ENDTIME_LEEWAY_MS before the start time of
+                            // its timeout window until TEARDOWN_LEEWAY_MS before the start time of
                             // the next scenario, before executing the scenario's @After methods.
                             // The above does not apply if current scenario is the last one, in
                             // which case the idle is never performed regardless of its after_test
@@ -96,8 +107,13 @@
                                     && mScenario
                                             .getAfterTest()
                                             .equals(Scenario.AfterTest.STAY_IN_APP)) {
+                                // Subtract the teardown leeway so that teardown methods can finish
+                                // within the scope of the timeout rule.
                                 performIdleBeforeTeardown(
-                                        max(getTimeRemaining() - ENDTIME_LEEWAY_MS, 0));
+                                        max(
+                                                getTimeRemainingForTimeoutRule()
+                                                        - TEARDOWN_LEEWAY_MS,
+                                                0));
                             }
                         }
                     }
@@ -128,19 +144,22 @@
                 InstrumentationRegistry.getInstrumentation(), mArguments);
         // If there are remaining scenarios, idle until the next one starts.
         if (mShouldIdle) {
-            performIdleBeforeNextScenario(getTimeRemaining());
+            performIdleBeforeNextScenario(getTimeRemainingForScenario());
         }
     }
 
-    /**
-     * Get the duration to idle after the current scenario. If the current scenario is the last one
-     * in the profile, returns 0.
-     */
-    private long getTimeRemaining() {
+    /** Get the remaining time within the current scenario. */
+    private long getTimeRemainingForScenario() {
         // The idle time is total time minus time elapsed since the current scenario started.
         return max(mTotalTimeoutMs - (System.currentTimeMillis() - mStartTimeMs), 0);
     }
 
+    /** Get the remaining time within the current timeout rule. */
+    private long getTimeRemainingForTimeoutRule() {
+        // The idle time is total time minus time elapsed since the current scenario started.
+        return max(mEnforcedTimeoutMs - (System.currentTimeMillis() - mStartTimeMs), 0);
+    }
+
     @VisibleForTesting
     protected void performIdleBeforeTeardown(long durationMs) {
         idleWithSystemClockSleep(durationMs);
diff --git a/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/ProfileSuiteTest.java b/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/ProfileSuiteTest.java
index 5f00b04..d42ccd7 100644
--- a/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/ProfileSuiteTest.java
+++ b/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/ProfileSuiteTest.java
@@ -288,7 +288,7 @@
                                     long expectedTimeout =
                                             suiteTimeoutMsecs
                                                     - TimeUnit.SECONDS.toMillis(4)
-                                                    - ScheduledScenarioRunner.ENDTIME_LEEWAY_MS;
+                                                    - ScheduledScenarioRunner.TEARDOWN_LEEWAY_MS;
                                     return abs(exceptionTimeout - expectedTimeout)
                                             <= SCHEDULE_LEEWAY_MS;
                                 });
diff --git a/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/ScheduledScenarioRunnerTest.java b/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/ScheduledScenarioRunnerTest.java
index 7c2e5fd..7afec46 100644
--- a/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/ScheduledScenarioRunnerTest.java
+++ b/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/ScheduledScenarioRunnerTest.java
@@ -53,7 +53,6 @@
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
-
 /** Unit tests for the {@link ScheduledScenarioRunner} runner. */
 @RunWith(JUnit4.class)
 public class ScheduledScenarioRunnerTest {
@@ -109,9 +108,9 @@
     @Test
     public void testOverTimeTest_throwsTestTimedOutException() throws InitializationError {
         ArgumentCaptor<Failure> failureCaptor = ArgumentCaptor.forClass(Failure.class);
-        // Set a over time test with a 5-second window that will idle until the end of the window is
+        // Set a over time test with a 6-second window that will idle until the end of the window is
         // reached.
-        long timeoutMs = TimeUnit.SECONDS.toMillis(5);
+        long timeoutMs = TimeUnit.SECONDS.toMillis(6);
         Scenario testScenario =
                 Scenario.newBuilder()
                         .setAt("00:00:00")
@@ -143,7 +142,7 @@
                                                     .getTimeUnit()
                                                     .toMillis(exception.getTimeout());
                                     long expectedTimeout =
-                                            timeoutMs - ScheduledScenarioRunner.ENDTIME_LEEWAY_MS;
+                                            timeoutMs - ScheduledScenarioRunner.TEARDOWN_LEEWAY_MS;
                                     return abs(exceptionTimeout - expectedTimeout)
                                             <= TIMEOUT_ERROR_MARGIN_MS;
                                 });
@@ -153,7 +152,7 @@
     /** Test that an over time test does not idle before teardown. */
     @Test
     public void testOverTimeTest_doesNotIdleBeforeTeardown() throws InitializationError {
-        // Set a over time test with a 5-second window that will idle until the end of the window is
+        // Set a over time test with a 6-second window that will idle until the end of the window is
         // reached.
         Scenario testScenario =
                 Scenario.newBuilder()
@@ -166,7 +165,7 @@
                         new ScheduledScenarioRunner(
                                 SampleProfileSuite.LongIdleTest.class,
                                 testScenario,
-                                TimeUnit.SECONDS.toMillis(5),
+                                TimeUnit.SECONDS.toMillis(6),
                                 true));
         runner.run(mRunNotifier);
         // There should not be idle before teardown as the test should not have left itself enough
@@ -177,7 +176,7 @@
     /** Test that an over time test still idles until tne next scenario is supposed to begin. */
     @Test
     public void testOverTimeTest_idlesAfterTeardownUntilNextScenario() throws InitializationError {
-        // Set a over time test with a 5-second window that will idle until the end of the window is
+        // Set a over time test with a 6-second window that will idle until the end of the window is
         // reached.
         Scenario testScenario =
                 Scenario.newBuilder()
@@ -190,7 +189,7 @@
                         new ScheduledScenarioRunner(
                                 SampleProfileSuite.LongIdleTest.class,
                                 testScenario,
-                                TimeUnit.SECONDS.toMillis(5),
+                                TimeUnit.SECONDS.toMillis(6),
                                 true));
         runner.run(mRunNotifier);
         // Verify that it still idles until the next scenario; duration should be roughly equal to
@@ -198,16 +197,16 @@
         verify(runner, times(1))
                 .performIdleBeforeNextScenario(
                         getWithinMarginMatcher(
-                                ScheduledScenarioRunner.ENDTIME_LEEWAY_MS,
+                                ScheduledScenarioRunner.TEARDOWN_LEEWAY_MS,
                                 TIMEOUT_ERROR_MARGIN_MS));
     }
 
     /** Test that a test set to stay in the app after the test idles after its @Test method. */
     @Test
     public void testRespectsAfterTestPolicy_stayInApp() throws InitializationError {
-        // Set a passing test with a 5-second timeout that will idle after its @Test method and
+        // Set a passing test with a 6-second timeout that will idle after its @Test method and
         // idle until the end of the timeout is reached.
-        long timeoutMs = TimeUnit.SECONDS.toMillis(5);
+        long timeoutMs = TimeUnit.SECONDS.toMillis(6);
         Scenario testScenario =
                 Scenario.newBuilder()
                         .setAt("00:00:00")
@@ -222,21 +221,24 @@
                                 timeoutMs,
                                 true));
         runner.run(mRunNotifier);
-        // Idles before teardown; duration should be roughly equal to the timeout minus the leeway
-        // set in {@link ScheduledScenarioRunner}.
+        // Idles before teardown; duration should be roughly equal to the timeout minus two times
+        // the leeway set in {@link ScheduledScenarioRunner}. Please see that class to see why two
+        // times the leeway is expected here.
         verify(runner, times(1))
                 .performIdleBeforeTeardown(
                         getWithinMarginMatcher(
-                                timeoutMs - ScheduledScenarioRunner.ENDTIME_LEEWAY_MS,
+                                timeoutMs - 2 * ScheduledScenarioRunner.TEARDOWN_LEEWAY_MS,
                                 TIMEOUT_ERROR_MARGIN_MS));
+        // Test should have passed.
+        verify(mRunNotifier, never()).fireTestFailure(any(Failure.class));
     }
 
     /** Test that a test set to exit the app after the test does not idle after its @Test method. */
     @Test
     public void testRespectsAfterTestPolicy_exit() throws InitializationError {
-        // Set a passing test with a 5-second timeout that does not idle after its @Test method and
+        // Set a passing test with a 6-second timeout that does not idle after its @Test method and
         // will idle until the end of the timeout is reached.
-        long timeoutMs = TimeUnit.SECONDS.toMillis(5);
+        long timeoutMs = TimeUnit.SECONDS.toMillis(6);
         Scenario testScenario =
                 Scenario.newBuilder()
                         .setAt("00:00:00")
@@ -257,12 +259,14 @@
         verify(runner, times(1))
                 .performIdleBeforeNextScenario(
                         getWithinMarginMatcher(timeoutMs, TIMEOUT_ERROR_MARGIN_MS));
+        // Test should have passed.
+        verify(mRunNotifier, never()).fireTestFailure(any(Failure.class));
     }
 
     /** Test that an ignored scenario still includes the timeout dictated in a profile. */
     @Test
     public void testIgnoredScenario_doesIdle() throws InitializationError, Exception {
-        long timeoutMs = TimeUnit.SECONDS.toMillis(5);
+        long timeoutMs = TimeUnit.SECONDS.toMillis(6);
         Scenario testScenario =
                 Scenario.newBuilder()
                         .setAt("00:00:00")
@@ -298,7 +302,7 @@
     /** Test that the last test does not have idle after it, regardless of its AfterTest policy. */
     @Test
     public void testLastScenarioDoesNotIdle() throws InitializationError {
-        // Set a passing test with a 5-second timeout that is set to idle after its @Test method and
+        // Set a passing test with a 6-second timeout that is set to idle after its @Test method and
         // but should not idle as it will be the last test in practice.
         Scenario testScenario =
                 Scenario.newBuilder()
@@ -311,7 +315,7 @@
                         new ScheduledScenarioRunner(
                                 SampleProfileSuite.PassingTest.class,
                                 testScenario,
-                                TimeUnit.SECONDS.toMillis(5),
+                                TimeUnit.SECONDS.toMillis(6),
                                 false));
         runner.run(mRunNotifier);
         // There should not be idle of any form.
@@ -338,7 +342,7 @@
                         new ScheduledScenarioRunner(
                                 ArgumentTest.class,
                                 testScenario,
-                                TimeUnit.SECONDS.toMillis(5),
+                                TimeUnit.SECONDS.toMillis(6),
                                 false));
         runner.run(mRunNotifier);
         verifyForAssertionFailures(mRunNotifier);
@@ -360,7 +364,7 @@
                         .build();
         ScheduledScenarioRunner runner =
                 new ScheduledScenarioRunner(
-                        ArgumentTest.class, testScenario, TimeUnit.SECONDS.toMillis(5), false);
+                        ArgumentTest.class, testScenario, TimeUnit.SECONDS.toMillis(6), false);
         runner.run(mRunNotifier);
         Bundle argsAfterTest = InstrumentationRegistry.getArguments();
         Assert.assertTrue(bundlesContainSameStringKeyValuePairs(argsBeforeTest, argsAfterTest));
diff --git a/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/samples/testing/SampleProfileSuite.java b/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/samples/testing/SampleProfileSuite.java
index 0756b4d..2b1bf10 100644
--- a/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/samples/testing/SampleProfileSuite.java
+++ b/libraries/health/runners/longevity/platform/tests/src/android/platform/test/longevity/samples/testing/SampleProfileSuite.java
@@ -20,6 +20,7 @@
 import android.platform.test.longevity.ProfileSuite;
 import android.platform.test.scenario.annotation.Scenario;
 
+import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -42,6 +43,12 @@
         public void testLongIdle() {
             SystemClock.sleep(TimeUnit.SECONDS.toMillis(10));
         }
+
+        // Simulates a quick teardown step.
+        @AfterClass
+        public static void dummyTearDown() {
+            SystemClock.sleep(100);
+        }
     }
 
     @Scenario
@@ -51,5 +58,11 @@
         public void testPassing() {
             Assert.assertEquals(1, 1);
         }
+
+        // Simulates a quick teardown step.
+        @AfterClass
+        public static void dummyTearDown() {
+            SystemClock.sleep(100);
+        }
     }
 }