Wait for device to charge between test cases am: 69f1b74d95
am: 2dc2a9c725

Change-Id: Idf4ff52bdca781f13a62c18c299060859dcf0768
diff --git a/src/com/android/nn/benchmark/app/BenchmarkTestBase.java b/src/com/android/nn/benchmark/app/BenchmarkTestBase.java
index 0b17d2a..5aaa056 100644
--- a/src/com/android/nn/benchmark/app/BenchmarkTestBase.java
+++ b/src/com/android/nn/benchmark/app/BenchmarkTestBase.java
@@ -18,6 +18,11 @@
 
 
 import android.app.Activity;
+import android.content.BroadcastReceiver;
+import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
+import android.os.BatteryManager;
 import android.os.Bundle;
 import android.os.Trace;
 import android.test.ActivityInstrumentationTestCase2;
@@ -87,6 +92,39 @@
         setUseNNApi(true);
     }
 
+    protected void waitUntilCharged() {
+        Log.v(NNBenchmark.TAG, "Waiting for the device to charge");
+
+        Object lock = new Object();
+        BroadcastReceiver receiver = new BroadcastReceiver() {
+            @Override
+            public void onReceive(Context context, Intent intent) {
+                int level = intent.getIntExtra(BatteryManager.EXTRA_LEVEL, -1);
+                int scale = intent.getIntExtra(BatteryManager.EXTRA_SCALE, -1);
+                int percentage  = level * 100 / scale;
+                Log.v(NNBenchmark.TAG, "Battery level: " + percentage + "%");
+
+                int status = intent.getIntExtra(BatteryManager.EXTRA_STATUS, -1);
+                if (status == BatteryManager.BATTERY_STATUS_FULL) {
+                    synchronized (lock) {
+                        lock.notify();
+                    }
+                } else if (status != BatteryManager.BATTERY_STATUS_CHARGING) {
+                    Log.e(NNBenchmark.TAG, "Device is not charging");
+                }
+            }
+        };
+
+        mActivity.registerReceiver(receiver, new IntentFilter(Intent.ACTION_BATTERY_CHANGED));
+        synchronized (lock) {
+            try {
+                lock.wait();
+            } catch (InterruptedException e) {
+            }
+        }
+        mActivity.unregisterReceiver(receiver);
+    }
+
     @Override
     @Before
     public void setUp() throws Exception {
diff --git a/src/com/android/nn/benchmark/app/NNInferenceStressTest.java b/src/com/android/nn/benchmark/app/NNInferenceStressTest.java
index c4fbdc1..9b4fdf4 100644
--- a/src/com/android/nn/benchmark/app/NNInferenceStressTest.java
+++ b/src/com/android/nn/benchmark/app/NNInferenceStressTest.java
@@ -66,6 +66,7 @@
     @Test
     @LargeTest
     public void stressTestNNAPI() throws IOException {
+        waitUntilCharged();
         setUseNNApi(true);
         setCompleteInputSet(false);
         TestAction ta = new TestAction(mModel, WARMUP_SECONDS, RUNTIME_SECONDS);
diff --git a/src/com/android/nn/benchmark/app/NNModelLoadingStressTest.java b/src/com/android/nn/benchmark/app/NNModelLoadingStressTest.java
index 1ab5e2e..faf507e 100644
--- a/src/com/android/nn/benchmark/app/NNModelLoadingStressTest.java
+++ b/src/com/android/nn/benchmark/app/NNModelLoadingStressTest.java
@@ -71,11 +71,12 @@
     @Test
     @LargeTest
     public void stressTestNNAPI() throws IOException {
-        Log.i(TAG, mModel.getTestName());
+        waitUntilCharged();
         setUseNNApi(true);
         setCompleteInputSet(true);
+        float endTime = stopwatch.runtime(TimeUnit.SECONDS) + RUNTIME_SECONDS;
         TestAction ta = new TestAction(mModel, WARMUP_SECONDS, INFERENCE_SECONDS);
-        while (stopwatch.runtime(TimeUnit.SECONDS) < RUNTIME_SECONDS) {
+        while (stopwatch.runtime(TimeUnit.SECONDS) < endTime) {
             runTest(ta, mModel.getTestName());
         }
     }