diff --git a/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkStateObserver.aidl b/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkStateObserver.aidl
index 09f3120..165f530 100644
--- a/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkStateObserver.aidl
+++ b/tests/cts/hostside/aidl/com/android/cts/net/hostside/INetworkStateObserver.aidl
@@ -16,6 +16,7 @@
 
 package com.android.cts.net.hostside;
 
-oneway interface INetworkStateObserver {
+interface INetworkStateObserver {
+    boolean isForeground();
     void onNetworkStateChecked(String resultData);
 }
\ No newline at end of file
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
index c5342ba..289de86 100644
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/AbstractRestrictBackgroundNetworkTestCase.java
@@ -31,12 +31,14 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
+import android.content.IntentFilter;
 import android.content.pm.PackageManager;
 import android.net.ConnectivityManager;
 import android.net.NetworkInfo;
 import android.net.NetworkInfo.DetailedState;
 import android.net.NetworkInfo.State;
 import android.net.wifi.WifiManager;
+import android.os.BatteryManager;
 import android.os.Binder;
 import android.os.Bundle;
 import android.os.SystemClock;
@@ -90,11 +92,19 @@
     protected static final int TYPE_COMPONENT_ACTIVTIY = 0;
     protected static final int TYPE_COMPONENT_FOREGROUND_SERVICE = 1;
 
+    private static final int BATTERY_STATE_TIMEOUT_MS = 5000;
+    private static final int BATTERY_STATE_CHECK_INTERVAL_MS = 500;
+
     private static final int FOREGROUND_PROC_NETWORK_TIMEOUT_MS = 6000;
 
     // Must be higher than NETWORK_TIMEOUT_MS
     private static final int ORDERED_BROADCAST_TIMEOUT_MS = NETWORK_TIMEOUT_MS * 4;
 
+    private static final IntentFilter BATTERY_CHANGED_FILTER =
+            new IntentFilter(Intent.ACTION_BATTERY_CHANGED);
+
+    private static final String APP_NOT_FOREGROUND_ERROR = "app_not_fg";
+
     protected Context mContext;
     protected Instrumentation mInstrumentation;
     protected ConnectivityManager mCm;
@@ -722,10 +732,29 @@
 
     protected void turnBatteryOff() throws Exception {
         executeSilentShellCommand("cmd battery unplug");
+        assertBatteryState(false);
     }
 
     protected void turnBatteryOn() throws Exception {
         executeSilentShellCommand("cmd battery reset");
+        assertBatteryState(true);
+
+    }
+
+    private void assertBatteryState(boolean pluggedIn) throws Exception {
+        final long endTime = SystemClock.elapsedRealtime() + BATTERY_STATE_TIMEOUT_MS;
+        while (isDevicePluggedIn() != pluggedIn && SystemClock.elapsedRealtime() <= 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 boolean isDevicePluggedIn() {
+        final Intent batteryIntent = mContext.registerReceiver(null, BATTERY_CHANGED_FILTER);
+        return batteryIntent.getIntExtra(BatteryManager.EXTRA_PLUGGED, -1) > 0;
     }
 
     protected void turnScreenOff() throws Exception {
@@ -795,11 +824,12 @@
     protected void registerBroadcastReceiver() throws Exception {
         mServiceClient.registerBroadcastReceiver();
 
+        final Intent intent = new Intent(ACTION_RECEIVER_READY)
+                .addFlags(Intent.FLAG_RECEIVER_FOREGROUND);
         // Wait until receiver is ready.
         final int maxTries = 10;
         for (int i = 1; i <= maxTries; i++) {
-            final String message =
-                    sendOrderedBroadcast(new Intent(ACTION_RECEIVER_READY), SECOND_IN_MS);
+            final String message = sendOrderedBroadcast(intent, SECOND_IN_MS * 4);
             Log.d(TAG, "app2 receiver acked: " + message);
             if (message != null) {
                 return;
@@ -859,7 +889,12 @@
             mContext.startActivity(launchIntent);
             if (latch.await(FOREGROUND_PROC_NETWORK_TIMEOUT_MS, TimeUnit.MILLISECONDS)) {
                 if (!errors[0].isEmpty()) {
-                    fail("Network is not available for app2 (" + mUid + "): " + errors[0]);
+                    if (errors[0] == APP_NOT_FOREGROUND_ERROR) {
+                        // App didn't come to foreground when the activity is started, so try again.
+                        assertForegroundNetworkAccess();
+                    } else {
+                        fail("Network is not available for app2 (" + mUid + "): " + errors[0]);
+                    }
                 }
             } else {
                 fail("Timed out waiting for network availability status from app2 (" + mUid + ")");
@@ -898,8 +933,21 @@
             final String[] errors) {
         return new INetworkStateObserver.Stub() {
             @Override
+            public boolean isForeground() {
+                try {
+                    final ProcessState state = getProcessStateByUid(mUid);
+                    return !isBackground(state.state);
+                } catch (Exception e) {
+                    Log.d(TAG, "Error while reading the proc state for " + mUid + ": " + e);
+                    return false;
+                }
+            }
+
+            @Override
             public void onNetworkStateChecked(String resultData) {
-                errors[0] = checkForAvailabilityInResultData(resultData, true);
+                errors[0] = resultData == null
+                        ? APP_NOT_FOREGROUND_ERROR
+                        : checkForAvailabilityInResultData(resultData, true);
                 latch.countDown();
             }
         };
diff --git a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/Common.java b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/Common.java
index 20bbd5a..351733e 100644
--- a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/Common.java
+++ b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/Common.java
@@ -72,12 +72,21 @@
         final INetworkStateObserver observer = INetworkStateObserver.Stub.asInterface(
                 extras.getBinder(KEY_NETWORK_STATE_OBSERVER));
         if (observer != null) {
+            try {
+                if (!observer.isForeground()) {
+                    Log.e(TAG, "App didn't come to foreground");
+                    observer.onNetworkStateChecked(null);
+                    return;
+                }
+            } catch (RemoteException e) {
+                Log.e(TAG, "Error occurred while reading the proc state: " + e);
+            }
             AsyncTask.execute(() -> {
                 try {
                     observer.onNetworkStateChecked(
                             MyBroadcastReceiver.checkNetworkStatus(context));
                 } catch (RemoteException e) {
-                    Log.e(TAG, "Error occured while notifying the observer: " + e);
+                    Log.e(TAG, "Error occurred while notifying the observer: " + e);
                 }
             });
         }
diff --git a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyActivity.java b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyActivity.java
index da7e704..286cc2f 100644
--- a/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyActivity.java
+++ b/tests/cts/hostside/app2/src/com/android/cts/net/hostside/app2/MyActivity.java
@@ -41,6 +41,7 @@
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
+        Log.d(TAG, "MyActivity.onCreate()");
         Common.notifyNetworkStateObserver(this, getIntent());
         finishCommandReceiver = new BroadcastReceiver() {
             @Override