diff --git a/tests/cts/hostside/app/AndroidManifest.xml b/tests/cts/hostside/app/AndroidManifest.xml
index d56e5d4..56d3cb5 100644
--- a/tests/cts/hostside/app/AndroidManifest.xml
+++ b/tests/cts/hostside/app/AndroidManifest.xml
@@ -30,6 +30,7 @@
     <uses-permission android:name="android.permission.WRITE_SECURE_SETTINGS"/>
     <uses-permission android:name="android.permission.WRITE_EXTERNAL_STORAGE"/>
     <uses-permission android:name="android.permission.MANAGE_EXTERNAL_STORAGE"/>
+    <uses-permission android:name="android.permission.WAKE_LOCK" />
 
     <application android:requestLegacyExternalStorage="true">
         <uses-library android:name="android.test.runner"/>
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 108a86e..a281aed 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
@@ -56,6 +56,7 @@
 import android.os.Bundle;
 import android.os.RemoteCallback;
 import android.os.SystemClock;
+import android.os.PowerManager;
 import android.provider.DeviceConfig;
 import android.service.notification.NotificationListenerService;
 import android.util.Log;
@@ -163,6 +164,8 @@
     private int mMyUid;
     private MyServiceClient mServiceClient;
     private DeviceConfigStateHelper mDeviceIdleDeviceConfigStateHelper;
+    private PowerManager mPowerManager;
+    private PowerManager.WakeLock mLock;
 
     @Rule
     public final RuleChain mRuleChain = RuleChain.outerRule(new RequiredPropertiesRule())
@@ -181,8 +184,10 @@
         mMyUid = getUid(mContext.getPackageName());
         mServiceClient = new MyServiceClient(mContext);
         mServiceClient.bind();
+        mPowerManager = mContext.getSystemService(PowerManager.class);
         executeShellCommand("cmd netpolicy start-watching " + mUid);
         setAppIdle(false);
+        mLock = mPowerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, TAG);
 
         Log.i(TAG, "Apps status:\n"
                 + "\ttest app: uid=" + mMyUid + ", state=" + getProcessStateByUid(mMyUid) + "\n"
@@ -192,6 +197,7 @@
     protected void tearDown() throws Exception {
         executeShellCommand("cmd netpolicy stop-watching");
         mServiceClient.unbind();
+        if (mLock.isHeld()) mLock.release();
     }
 
     protected int getUid(String packageName) throws Exception {
@@ -695,11 +701,13 @@
     }
 
     protected void turnScreenOff() throws Exception {
+        if (!mLock.isHeld()) mLock.acquire();
         executeSilentShellCommand("input keyevent KEYCODE_SLEEP");
     }
 
     protected void turnScreenOn() throws Exception {
         executeSilentShellCommand("input keyevent KEYCODE_WAKEUP");
+        if (mLock.isHeld()) mLock.release();
         executeSilentShellCommand("wm dismiss-keyguard");
     }