/*
 * Copyright (C) 2017 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.android.internal.os;

import static android.os.BatteryStats.UID_TIMES_TYPE_ALL;
import static android.os.BatteryStats.Uid.NUM_PROCESS_STATE;
import static android.os.BatteryStats.Uid.PROCESS_STATE_BACKGROUND;
import static android.os.BatteryStats.Uid.PROCESS_STATE_CACHED;
import static android.os.BatteryStats.Uid.PROCESS_STATE_FOREGROUND;
import static android.os.BatteryStats.Uid.PROCESS_STATE_FOREGROUND_SERVICE;
import static android.os.BatteryStats.Uid.PROCESS_STATE_TOP;
import static android.os.BatteryStats.Uid.PROCESS_STATE_TOP_SLEEPING;
import static android.os.BatteryStats.Uid.UID_PROCESS_TYPES;

import static junit.framework.Assert.assertNotNull;
import static junit.framework.Assert.assertNull;
import static junit.framework.Assert.assertTrue;
import static junit.framework.Assert.fail;

import com.android.frameworks.coretests.aidl.ICmdCallback;
import com.android.frameworks.coretests.aidl.ICmdReceiver;

import android.app.ActivityManager;
import android.app.KeyguardManager;
import android.content.ComponentName;
import android.content.Context;
import android.content.Intent;
import android.content.IntentFilter;
import android.content.ServiceConnection;
import android.content.pm.PackageManager;
import android.os.BatteryManager;
import android.os.BatteryStats;
import android.os.Bundle;
import android.os.IBinder;
import android.os.PowerManager;
import android.os.Process;
import android.os.SystemClock;
import android.support.test.InstrumentationRegistry;
import android.support.test.filters.LargeTest;
import android.support.test.runner.AndroidJUnit4;
import android.support.test.uiautomator.UiDevice;
import android.util.DebugUtils;
import android.util.Log;

import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;

import java.util.Arrays;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

@LargeTest
@RunWith(AndroidJUnit4.class)
public class BstatsCpuTimesValidationTest {
    private static final String TAG = BstatsCpuTimesValidationTest.class.getSimpleName();

    private static final String TEST_PKG = "com.android.coretests.apps.bstatstestapp";
    private static final String TEST_ACTIVITY = TEST_PKG + ".TestActivity";
    private static final String TEST_SERVICE = TEST_PKG + ".TestService";
    private static final String ISOLATED_TEST_SERVICE = TEST_PKG + ".IsolatedTestService";

    private static final String EXTRA_KEY_CMD_RECEIVER = "cmd_receiver";
    private static final int FLAG_START_FOREGROUND = 1;

    private static final int BATTERY_STATE_TIMEOUT_MS = 2000;
    private static final int BATTERY_STATE_CHECK_INTERVAL_MS = 200;

    private static final int START_ACTIVITY_TIMEOUT_MS = 2000;
    private static final int START_FG_SERVICE_TIMEOUT_MS = 2000;
    private static final int START_SERVICE_TIMEOUT_MS = 2000;
    private static final int START_ISOLATED_SERVICE_TIMEOUT_MS = 2000;

    private static final int GENERAL_TIMEOUT_MS = 4000;
    private static final int GENERAL_INTERVAL_MS = 200;

    private static final int WORK_DURATION_MS = 2000;

    private static Context sContext;
    private static UiDevice sUiDevice;
    private static int sTestPkgUid;
    private static boolean sCpuFreqTimesAvailable;
    private static boolean sPerProcStateTimesAvailable;

    @BeforeClass
    public static void setupOnce() throws Exception {
        sContext = InstrumentationRegistry.getContext();
        sUiDevice = UiDevice.getInstance(InstrumentationRegistry.getInstrumentation());
        sContext.getPackageManager().setApplicationEnabledSetting(TEST_PKG,
                PackageManager.COMPONENT_ENABLED_STATE_ENABLED, 0);
        sTestPkgUid = sContext.getPackageManager().getPackageUid(TEST_PKG, 0);
        checkCpuTimesAvailability();
    }

    // Checks cpu freq times of system uid as an indication of whether /proc/uid_time_in_state
    // and /proc/uid/<uid>/time_in_state kernel nodes are available.
    private static void checkCpuTimesAvailability() throws Exception {
        batteryOn();
        SystemClock.sleep(GENERAL_TIMEOUT_MS);
        batteryOff();
        final long[] totalCpuTimes = getAllCpuFreqTimes(Process.SYSTEM_UID);
        sCpuFreqTimesAvailable = totalCpuTimes != null;
        final long[] fgSvcCpuTimes = getAllCpuFreqTimes(Process.SYSTEM_UID,
                PROCESS_STATE_FOREGROUND_SERVICE);
        sPerProcStateTimesAvailable = fgSvcCpuTimes != null;
    }

    @Test
    public void testCpuFreqTimes() throws Exception {
        if (!sCpuFreqTimesAvailable) {
            return;
        }

        batteryOnScreenOn();
        forceStop();
        resetBatteryStats();
        final long[] initialSnapshot = getAllCpuFreqTimes(sTestPkgUid);
        assertNull("Initial snapshot should be null, initial=" + Arrays.toString(initialSnapshot),
                initialSnapshot);
        doSomeWork();
        forceStop();

        final long[] cpuTimesMs = getAllCpuFreqTimes(sTestPkgUid);
        assertCpuTimesValid(cpuTimesMs);
        long actualCpuTimeMs = 0;
        for (int i = 0; i < cpuTimesMs.length / 2; ++i) {
            actualCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect total cpu time", WORK_DURATION_MS, actualCpuTimeMs);
        batteryOffScreenOn();
    }

    @Test
    public void testCpuFreqTimes_screenOff() throws Exception {
        if (!sCpuFreqTimesAvailable) {
            return;
        }

        batteryOnScreenOff();
        forceStop();
        resetBatteryStats();
        final long[] initialSnapshot = getAllCpuFreqTimes(sTestPkgUid);
        assertNull("Initial snapshot should be null, initial=" + Arrays.toString(initialSnapshot),
                initialSnapshot);
        doSomeWork();
        forceStop();

        final long[] cpuTimesMs = getAllCpuFreqTimes(sTestPkgUid);
        assertCpuTimesValid(cpuTimesMs);
        long actualTotalCpuTimeMs = 0;
        for (int i = 0; i < cpuTimesMs.length / 2; ++i) {
            actualTotalCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect total cpu time", WORK_DURATION_MS, actualTotalCpuTimeMs);
        long actualScreenOffCpuTimeMs = 0;
        for (int i = cpuTimesMs.length / 2; i < cpuTimesMs.length; ++i) {
            actualScreenOffCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect screen-off cpu time",
                WORK_DURATION_MS, actualScreenOffCpuTimeMs);
        batteryOffScreenOn();
    }

    @Test
    public void testCpuFreqTimes_isolatedProcess() throws Exception {
        if (!sCpuFreqTimesAvailable) {
            return;
        }

        batteryOnScreenOn();
        forceStop();
        resetBatteryStats();
        final long[] initialSnapshot = getAllCpuFreqTimes(sTestPkgUid);
        assertNull("Initial snapshot should be null, initial=" + Arrays.toString(initialSnapshot),
                initialSnapshot);
        doSomeWorkInIsolatedProcess();
        forceStop();

        final long[] cpuTimesMs = getAllCpuFreqTimes(sTestPkgUid);
        assertCpuTimesValid(cpuTimesMs);
        long actualCpuTimeMs = 0;
        for (int i = 0; i < cpuTimesMs.length / 2; ++i) {
            actualCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect total cpu time", WORK_DURATION_MS, actualCpuTimeMs);
        batteryOffScreenOn();
    }

    @Test
    public void testCpuFreqTimes_stateTop() throws Exception {
        if (!sCpuFreqTimesAvailable || !sPerProcStateTimesAvailable) {
            return;
        }

        batteryOnScreenOn();
        forceStop();
        resetBatteryStats();
        final long[] initialSnapshot = getAllCpuFreqTimes(sTestPkgUid);
        assertNull("Initial snapshot should be null, initial=" + Arrays.toString(initialSnapshot),
                initialSnapshot);
        assertNull("Initial top state snapshot should be null",
                getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_TOP));

        doSomeWork(PROCESS_STATE_TOP);
        forceStop();

        final long[] cpuTimesMs = getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_TOP);
        final String msgCpuTimes = getAllCpuTimesMsg();
        assertCpuTimesValid(cpuTimesMs);
        long actualCpuTimeMs = 0;
        for (int i = 0; i < cpuTimesMs.length / 2; ++i) {
            actualCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect total cpu time, " + msgCpuTimes,
                WORK_DURATION_MS, actualCpuTimeMs);
        batteryOffScreenOn();
    }

    @Test
    public void testIsolatedCpuFreqTimes_stateService() throws Exception {
        if (!sCpuFreqTimesAvailable || !sPerProcStateTimesAvailable) {
            return;
        }

        batteryOnScreenOn();
        forceStop();
        resetBatteryStats();
        final long[] initialSnapshot = getAllCpuFreqTimes(sTestPkgUid);
        assertNull("Initial snapshot should be null, initial=" + Arrays.toString(initialSnapshot),
                initialSnapshot);
        assertNull("Initial top state snapshot should be null",
                getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_TOP));

        final ICmdReceiver activityReceiver = ICmdReceiver.Stub.asInterface(startActivity());
        final ICmdReceiver isolatedReceiver = ICmdReceiver.Stub.asInterface(startIsolatedService());
        try {
            assertProcState(PROCESS_STATE_TOP);
            isolatedReceiver.doSomeWork(WORK_DURATION_MS);
        } finally {
            activityReceiver.finishHost();
            isolatedReceiver.finishHost();
        }
        forceStop();

        final long[] cpuTimesMs = getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_TOP);
        final String msgCpuTimes = getAllCpuTimesMsg();
        assertCpuTimesValid(cpuTimesMs);
        long actualCpuTimeMs = 0;
        for (int i = 0; i < cpuTimesMs.length / 2; ++i) {
            actualCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect total cpu time, " + msgCpuTimes,
                WORK_DURATION_MS, actualCpuTimeMs);
        batteryOffScreenOn();
    }

    @Test
    public void testCpuFreqTimes_stateTopSleeping() throws Exception {
        if (!sCpuFreqTimesAvailable || !sPerProcStateTimesAvailable) {
            return;
        }

        batteryOnScreenOff();
        forceStop();
        resetBatteryStats();
        final long[] initialSnapshot = getAllCpuFreqTimes(sTestPkgUid);
        assertNull("Initial snapshot should be null, initial=" + Arrays.toString(initialSnapshot),
                initialSnapshot);
        assertNull("Initial top state snapshot should be null",
                getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_TOP_SLEEPING));

        doSomeWork(PROCESS_STATE_TOP_SLEEPING);
        forceStop();

        final long[] cpuTimesMs = getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_TOP_SLEEPING);
        final String msgCpuTimes = getAllCpuTimesMsg();
        assertCpuTimesValid(cpuTimesMs);
        long actualCpuTimeMs = 0;
        for (int i = cpuTimesMs.length / 2; i < cpuTimesMs.length; ++i) {
            actualCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect total cpu time, " + msgCpuTimes,
                WORK_DURATION_MS, actualCpuTimeMs);
        batteryOffScreenOn();
    }

    @Test
    public void testCpuFreqTimes_stateFgService() throws Exception {
        if (!sCpuFreqTimesAvailable || !sPerProcStateTimesAvailable) {
            return;
        }

        batteryOnScreenOff();
        forceStop();
        resetBatteryStats();
        final long[] initialSnapshot = getAllCpuFreqTimes(sTestPkgUid);
        assertNull("Initial snapshot should be null, initial=" + Arrays.toString(initialSnapshot),
                initialSnapshot);
        assertNull("Initial top state snapshot should be null",
                getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_FOREGROUND_SERVICE));

        doSomeWork(PROCESS_STATE_FOREGROUND_SERVICE);
        forceStop();

        final long[] cpuTimesMs = getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_FOREGROUND_SERVICE);
        final String msgCpuTimes = getAllCpuTimesMsg();
        assertCpuTimesValid(cpuTimesMs);
        long actualCpuTimeMs = 0;
        for (int i = 0; i < cpuTimesMs.length / 2; ++i) {
            actualCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect total cpu time, " + msgCpuTimes,
                WORK_DURATION_MS, actualCpuTimeMs);
        batteryOffScreenOn();
    }

    @Test
    public void testCpuFreqTimes_stateFg() throws Exception {
        if (!sCpuFreqTimesAvailable || !sPerProcStateTimesAvailable) {
            return;
        }

        batteryOnScreenOn();
        forceStop();
        resetBatteryStats();
        final long[] initialSnapshot = getAllCpuFreqTimes(sTestPkgUid);
        assertNull("Initial snapshot should be null, initial=" + Arrays.toString(initialSnapshot),
                initialSnapshot);
        assertNull("Initial top state snapshot should be null",
                getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_FOREGROUND));

        doSomeWork(PROCESS_STATE_FOREGROUND);
        forceStop();

        final long[] cpuTimesMs = getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_FOREGROUND);
        final String msgCpuTimes = getAllCpuTimesMsg();
        assertCpuTimesValid(cpuTimesMs);
        long actualCpuTimeMs = 0;
        for (int i = 0; i < cpuTimesMs.length / 2; ++i) {
            actualCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect total cpu time, " + msgCpuTimes,
                WORK_DURATION_MS, actualCpuTimeMs);
        batteryOff();
    }

    @Test
    public void testCpuFreqTimes_stateBg() throws Exception {
        if (!sCpuFreqTimesAvailable || !sPerProcStateTimesAvailable) {
            return;
        }

        batteryOnScreenOff();
        forceStop();
        resetBatteryStats();
        final long[] initialSnapshot = getAllCpuFreqTimes(sTestPkgUid);
        assertNull("Initial snapshot should be null, initial=" + Arrays.toString(initialSnapshot),
                initialSnapshot);
        assertNull("Initial top state snapshot should be null",
                getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_BACKGROUND));

        doSomeWork(PROCESS_STATE_BACKGROUND);
        forceStop();

        final long[] cpuTimesMs = getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_BACKGROUND);
        final String msgCpuTimes = getAllCpuTimesMsg();
        assertCpuTimesValid(cpuTimesMs);
        long actualCpuTimeMs = 0;
        for (int i = 0; i < cpuTimesMs.length / 2; ++i) {
            actualCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect total cpu time, " + msgCpuTimes,
                WORK_DURATION_MS, actualCpuTimeMs);
        batteryOffScreenOn();
    }

    @Test
    public void testCpuFreqTimes_stateCached() throws Exception {
        if (!sCpuFreqTimesAvailable || !sPerProcStateTimesAvailable) {
            return;
        }

        batteryOnScreenOn();
        forceStop();
        resetBatteryStats();
        final long[] initialSnapshot = getAllCpuFreqTimes(sTestPkgUid);
        assertNull("Initial snapshot should be null, initial=" + Arrays.toString(initialSnapshot),
                initialSnapshot);
        assertNull("Initial top state snapshot should be null",
                getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_CACHED));

        doSomeWork(PROCESS_STATE_CACHED);
        forceStop();

        final long[] cpuTimesMs = getAllCpuFreqTimes(sTestPkgUid, PROCESS_STATE_CACHED);
        final String msgCpuTimes = getAllCpuTimesMsg();
        assertCpuTimesValid(cpuTimesMs);
        long actualCpuTimeMs = 0;
        for (int i = 0; i < cpuTimesMs.length / 2; ++i) {
            actualCpuTimeMs += cpuTimesMs[i];
        }
        assertApproximateValue("Incorrect total cpu time, " + msgCpuTimes,
                WORK_DURATION_MS, actualCpuTimeMs);
        batteryOffScreenOn();
    }

    private void assertCpuTimesValid(long[] cpuTimes) {
        assertNotNull(cpuTimes);
        for (int i = 0; i < cpuTimes.length; ++i) {
            if (cpuTimes[i] < 0) {
                fail("Malformed cpu times data (-ve values): " + Arrays.toString(cpuTimes));
            }
        }
        final int numFreqs = cpuTimes.length / 2;
        for (int i = 0; i < numFreqs; ++i) {
            if (cpuTimes[i] < cpuTimes[numFreqs + i]) {
                fail("Malformed cpu times data (screen-off > total)" + Arrays.toString(cpuTimes));
            }
        }
    }

    private void assertApproximateValue(String errorPrefix, long expectedValue, long actualValue) {
        assertValueRange(errorPrefix, actualValue, expectedValue * 0.5, expectedValue * 1.5);
    }

    private void assertValueRange(String errorPrefix,
            long actualvalue, double minValue, double maxValue) {
        final String errorMsg = String.format(errorPrefix + "; actual=%s; min=%s; max=%s",
                actualvalue, minValue, maxValue);
        assertTrue(errorMsg, actualvalue < maxValue);
        assertTrue(errorMsg, actualvalue > minValue);
    }

    private void doSomeWork() throws Exception {
        final ICmdReceiver receiver = ICmdReceiver.Stub.asInterface(startActivity());
        receiver.doSomeWork(WORK_DURATION_MS);
        receiver.finishHost();
    }

    private void doSomeWorkInIsolatedProcess() throws Exception {
        final ICmdReceiver receiver = ICmdReceiver.Stub.asInterface(startIsolatedService());
        receiver.doSomeWork(WORK_DURATION_MS);
        receiver.finishHost();
    }

    private void doSomeWork(int procState) throws Exception {
        final ICmdReceiver receiver;
        switch (procState) {
            case PROCESS_STATE_TOP:
                receiver = ICmdReceiver.Stub.asInterface(startActivity());
                break;
            case PROCESS_STATE_TOP_SLEEPING:
                receiver = ICmdReceiver.Stub.asInterface(startActivity());
                break;
            case PROCESS_STATE_FOREGROUND_SERVICE:
                receiver = ICmdReceiver.Stub.asInterface(startForegroundService());
                break;
            case PROCESS_STATE_FOREGROUND:
                receiver = ICmdReceiver.Stub.asInterface(startService());
                receiver.showApplicationOverlay();
                break;
            case PROCESS_STATE_BACKGROUND:
                receiver = ICmdReceiver.Stub.asInterface(startService());
                break;
            case PROCESS_STATE_CACHED:
                receiver = ICmdReceiver.Stub.asInterface(startActivity());
                receiver.finishHost();
                break;
            default:
                throw new IllegalArgumentException("Unknown state: " + procState);
        }
        try {
            assertProcState(procState);
            receiver.doSomeWork(WORK_DURATION_MS);
        } finally {
            receiver.finishHost();
        }
    }

    private void assertProcState(String state) throws Exception {
        final String expectedState = "(" + state + ")";
        assertDelayedCondition("", () -> {
            final String uidStateStr = executeCmd("cmd activity get-uid-state " + sTestPkgUid);
            final String actualState = uidStateStr.split(" ")[1];
            return expectedState.equals(actualState) ? null
                    : "expected=" + expectedState + ", actual" + actualState;
        });
    }

    private void assertProcState(int expectedState) throws Exception {
        assertDelayedCondition("Unexpected proc state", () -> {
            final String uidStateStr = executeCmd("cmd activity get-uid-state " + sTestPkgUid);
            final int amProcState = Integer.parseInt(uidStateStr.split(" ")[0]);
            final int actualState = BatteryStats.mapToInternalProcessState(amProcState);
            return (actualState == expectedState) ? null
                    : "expected=" + getStateName(BatteryStats.Uid.class, expectedState)
                            + ", actual=" + getStateName(BatteryStats.Uid.class, actualState)
                            + ", amState=" + getStateName(ActivityManager.class, amProcState);
        });
    }

    private String getStateName(Class clazz, int procState) {
        return DebugUtils.valueToString(clazz, "PROCESS_STATE_", procState);
    }

    private IBinder startIsolatedService() throws Exception {
        final CountDownLatch latch = new CountDownLatch(1);
        final IBinder[] binders = new IBinder[1];
        final ServiceConnection connection = new ServiceConnection() {
            @Override
            public void onServiceConnected(ComponentName name, IBinder service) {
                binders[0] = service;
                latch.countDown();
            }

            @Override
            public void onServiceDisconnected(ComponentName name) {
            }
        };
        final Intent launchIntent = new Intent()
                .setComponent(new ComponentName(TEST_PKG, ISOLATED_TEST_SERVICE));
        sContext.bindService(launchIntent, connection, Context.BIND_AUTO_CREATE
                | Context.BIND_ALLOW_OOM_MANAGEMENT | Context.BIND_NOT_FOREGROUND);
        if (latch.await(START_ISOLATED_SERVICE_TIMEOUT_MS, TimeUnit.MILLISECONDS)) {
            if (binders[0] == null) {
                fail("Receiver binder should not be null");
            }
            return binders[0];
        } else {
            fail("Timed out waiting for the isolated test service to start");
        }
        return null;
    }

    private IBinder startForegroundService() throws Exception {
        final CountDownLatch latch = new CountDownLatch(1);
        final Intent launchIntent = new Intent()
                .setComponent(new ComponentName(TEST_PKG, TEST_SERVICE))
                .setFlags(FLAG_START_FOREGROUND);
        final Bundle extras = new Bundle();
        final IBinder[] binders = new IBinder[1];
        extras.putBinder(EXTRA_KEY_CMD_RECEIVER, new ICmdCallback.Stub() {
            @Override
            public void onLaunched(IBinder receiver) {
                binders[0] = receiver;
                latch.countDown();
            }
        });
        launchIntent.putExtras(extras);
        sContext.startForegroundService(launchIntent);
        if (latch.await(START_FG_SERVICE_TIMEOUT_MS, TimeUnit.MILLISECONDS)) {
            if (binders[0] == null) {
                fail("Receiver binder should not be null");
            }
            return binders[0];
        } else {
            fail("Timed out waiting for the test fg service to start; testUid=" + sTestPkgUid);
        }
        return null;
    }

    private IBinder startService() throws Exception {
        final CountDownLatch latch = new CountDownLatch(1);
        final Intent launchIntent = new Intent()
                .setComponent(new ComponentName(TEST_PKG, TEST_SERVICE));
        final Bundle extras = new Bundle();
        final IBinder[] binders = new IBinder[1];
        extras.putBinder(EXTRA_KEY_CMD_RECEIVER, new ICmdCallback.Stub() {
            @Override
            public void onLaunched(IBinder receiver) {
                binders[0] = receiver;
                latch.countDown();
            }
        });
        launchIntent.putExtras(extras);
        sContext.startService(launchIntent);
        if (latch.await(START_SERVICE_TIMEOUT_MS, TimeUnit.MILLISECONDS)) {
            if (binders[0] == null) {
                fail("Receiver binder should not be null");
            }
            return binders[0];
        } else {
            fail("Timed out waiting for the test service to start; testUid=" + sTestPkgUid);
        }
        return null;
    }

    private IBinder startActivity() throws Exception {
        final CountDownLatch latch = new CountDownLatch(1);
        final Intent launchIntent = new Intent()
                .setComponent(new ComponentName(TEST_PKG, TEST_ACTIVITY));
        final Bundle extras = new Bundle();
        final IBinder[] binders = new IBinder[1];
        extras.putBinder(EXTRA_KEY_CMD_RECEIVER, new ICmdCallback.Stub() {
            @Override
            public void onLaunched(IBinder receiver) {
                binders[0] = receiver;
                latch.countDown();
            }
        });
        launchIntent.putExtras(extras);
        sContext.startActivity(launchIntent);
        if (latch.await(START_ACTIVITY_TIMEOUT_MS, TimeUnit.MILLISECONDS)) {
            if (binders[0] == null) {
                fail("Receiver binder should not be null");
            }
            return binders[0];
        } else {
            fail("Timed out waiting for the test activity to start; testUid=" + sTestPkgUid);
        }
        return null;
    }

    private static String getAllCpuTimesMsg() throws Exception {
        final StringBuilder sb = new StringBuilder();
        sb.append("uid=" + sTestPkgUid + ";");
        sb.append(UID_TIMES_TYPE_ALL + "=" + getMsgCpuTimesSum(getAllCpuFreqTimes(sTestPkgUid)));
        for (int i = 0; i < NUM_PROCESS_STATE; ++i) {
            sb.append("|");
            sb.append(UID_PROCESS_TYPES[i] + "="
                    + getMsgCpuTimesSum(getAllCpuFreqTimes(sTestPkgUid, i)));
        }
        return sb.toString();
    }

    private static String getMsgCpuTimesSum(long[] cpuTimes) throws Exception {
        if (cpuTimes == null) {
            return "(0,0)";
        }
        long totalTime = 0;
        for (int i = 0; i < cpuTimes.length / 2; ++i) {
            totalTime += cpuTimes[i];
        }
        long screenOffTime = 0;
        for (int i = cpuTimes.length / 2; i < cpuTimes.length; ++i) {
            screenOffTime += cpuTimes[i];
        }
        return "(" + totalTime + "," + screenOffTime + ")";
    }

    private static long[] getAllCpuFreqTimes(int uid) throws Exception {
        final String checkinDump = executeCmdSilent("dumpsys batterystats --checkin");
        final Pattern pattern = Pattern.compile(uid + ",l,ctf," + UID_TIMES_TYPE_ALL + ",(.*?)\n");
        final Matcher matcher = pattern.matcher(checkinDump);
        if (!matcher.find()) {
            return null;
        }
        return parseCpuTimesStr(matcher.group(1));
    }

    private static long[] getAllCpuFreqTimes(int uid, int procState) throws Exception {
        final String checkinDump = executeCmdSilent("dumpsys batterystats --checkin");
        final Pattern pattern = Pattern.compile(
                uid + ",l,ctf," + UID_PROCESS_TYPES[procState] + ",(.*?)\n");
        final Matcher matcher = pattern.matcher(checkinDump);
        if (!matcher.find()) {
            return null;
        }
        return parseCpuTimesStr(matcher.group(1));
    }

    private static long[] parseCpuTimesStr(String str) {
        final String[] cpuTimesStr = str.split(",");
        final int freqCount = Integer.parseInt(cpuTimesStr[0]);
        if (cpuTimesStr.length != (2 * freqCount + 1)) {
            fail("Malformed data: " + Arrays.toString(cpuTimesStr));
        }
        final long[] cpuTimes = new long[freqCount * 2];
        for (int i = 0; i < cpuTimes.length; ++i) {
            cpuTimes[i] = Long.parseLong(cpuTimesStr[i + 1]);
        }
        return cpuTimes;
    }

    private void resetBatteryStats() throws Exception {
        executeCmd("dumpsys batterystats --reset");
    }

    private void batteryOnScreenOn() throws Exception {
        batteryOn();
        screenOn();
    }

    private void batteryOnScreenOff() throws Exception {
        batteryOn();
        screenoff();
    }

    private void batteryOffScreenOn() throws Exception {
        batteryOff();
        screenOn();
    }

    private static void batteryOn() throws Exception {
        executeCmd("dumpsys battery unplug");
        assertBatteryState(false);
    }

    private static void batteryOff() throws Exception {
        executeCmd("dumpsys battery reset");
        assertBatteryState(true);
    }

    private void screenOn() throws Exception {
        executeCmd("input keyevent KEYCODE_WAKEUP");
        executeCmd("wm dismiss-keyguard");
        assertKeyguardUnLocked();
        assertScreenInteractive(true);
    }

    private void screenoff() throws Exception {
        executeCmd("input keyevent KEYCODE_SLEEP");
        assertScreenInteractive(false);
    }

    private void forceStop() throws Exception {
        executeCmd("cmd activity force-stop " + TEST_PKG);
        assertProcState("NONEXISTENT");
    }

    private void assertKeyguardUnLocked() throws Exception {
        final KeyguardManager keyguardManager =
                (KeyguardManager) sContext.getSystemService(Context.KEYGUARD_SERVICE);
        assertDelayedCondition("Unexpected Keyguard state", () ->
                keyguardManager.isKeyguardLocked() ? "expected=unlocked" : null
        );
    }

    private void assertScreenInteractive(boolean interactive) throws Exception {
        final PowerManager powerManager =
                (PowerManager) sContext.getSystemService(Context.POWER_SERVICE);
        assertDelayedCondition("Unexpected screen interactive state", () ->
                interactive == powerManager.isInteractive() ? null : "expected=" + interactive
        );
    }

    private void assertDelayedCondition(String errMsgPrefix, ExpectedCondition condition)
            throws Exception {
        final long endTime = SystemClock.uptimeMillis() + GENERAL_TIMEOUT_MS;
        while (SystemClock.uptimeMillis() <= endTime) {
            if (condition.getErrIfNotTrue() == null) {
                return;
            }
            SystemClock.sleep(GENERAL_INTERVAL_MS);
        }
        final String errMsg = condition.getErrIfNotTrue();
        if (errMsg != null) {
            fail(errMsgPrefix + ": " + errMsg);
        }
    }

    private static void assertBatteryState(boolean pluggedIn) throws Exception {
        final long endTime = SystemClock.uptimeMillis() + BATTERY_STATE_TIMEOUT_MS;
        while (isDevicePluggedIn() != pluggedIn && SystemClock.uptimeMillis() <= endTime) {
            Thread.sleep(BATTERY_STATE_CHECK_INTERVAL_MS);
        }
        if (isDevicePluggedIn() != pluggedIn) {
            fail("Timed out waiting for the plugged-in state to change,"
                    + " expected pluggedIn: " + pluggedIn);
        }
    }

    private static boolean isDevicePluggedIn() {
        final Intent batteryIntent = sContext.registerReceiver(null,
                new IntentFilter(Intent.ACTION_BATTERY_CHANGED));
        return batteryIntent.getIntExtra(BatteryManager.EXTRA_PLUGGED, -1) > 0;
    }

    private static String executeCmd(String cmd) throws Exception {
        final String result = sUiDevice.executeShellCommand(cmd).trim();
        Log.d(TAG, String.format("Result for '%s': %s", cmd, result));
        return result;
    }

    private static String executeCmdSilent(String cmd) throws Exception {
        return sUiDevice.executeShellCommand(cmd).trim();
    }

    private interface ExpectedCondition {
        String getErrIfNotTrue() throws Exception;
    }
}