Merge "CTS: WifiLockState"
diff --git a/hostsidetests/statsd/apps/statsdapp/src/com/android/server/cts/device/statsd/AtomTests.java b/hostsidetests/statsd/apps/statsdapp/src/com/android/server/cts/device/statsd/AtomTests.java
index 825a2c4..0f1807e 100644
--- a/hostsidetests/statsd/apps/statsdapp/src/com/android/server/cts/device/statsd/AtomTests.java
+++ b/hostsidetests/statsd/apps/statsdapp/src/com/android/server/cts/device/statsd/AtomTests.java
@@ -246,6 +246,16 @@
     }
 
     @Test
+    public void testWifiLock() {
+        Context context = InstrumentationRegistry.getContext();
+        WifiManager wm = context.getSystemService(WifiManager.class);
+        WifiManager.WifiLock lock = wm.createWifiLock("StatsdCTSWifiLock");
+        lock.acquire();
+        sleep(500);
+        lock.release();
+    }
+
+    @Test
     /** Does two wifi scans. */
     // TODO: Copied this from BatterystatsValidation but we probably don't need to wait for results.
     public void testWifiScan() {
diff --git a/hostsidetests/statsd/src/android/cts/statsd/UidAtomTests.java b/hostsidetests/statsd/src/android/cts/statsd/UidAtomTests.java
index c02e8bd..b49369c 100644
--- a/hostsidetests/statsd/src/android/cts/statsd/UidAtomTests.java
+++ b/hostsidetests/statsd/src/android/cts/statsd/UidAtomTests.java
@@ -24,6 +24,7 @@
 import com.android.os.AtomsProto.FlashlightStateChanged;
 import com.android.os.AtomsProto.GpsScanStateChanged;
 import com.android.os.AtomsProto.WakeupAlarmOccurred;
+import com.android.os.AtomsProto.WifiLockStateChanged;
 import com.android.os.AtomsProto.WifiScanStateChanged;
 import com.android.os.StatsLog.EventMetricData;
 import com.android.tradefed.log.LogUtil;
@@ -209,6 +210,28 @@
         }
     }
 
+    public void testWifiLock() throws Exception {
+        if (!TESTS_ENABLED) return;
+        if (!hasFeature(FEATURE_WIFI, true)) return;
+
+        final int atomTag = Atom.WIFI_LOCK_STATE_CHANGED_FIELD_NUMBER;
+        Set<Integer> lockOn = new HashSet<>(Arrays.asList(WifiLockStateChanged.State.ON_VALUE));
+        Set<Integer> lockOff = new HashSet<>(Arrays.asList(WifiLockStateChanged.State.OFF_VALUE));
+
+        // Add state sets to the list in order.
+        List<Set<Integer>> stateSet = Arrays.asList(lockOn, lockOff);
+
+        createAndUploadConfig(atomTag);
+        runDeviceTests(DEVICE_SIDE_TEST_PACKAGE, ".AtomTests", "testWifiLock");
+
+        // Sorted list of events in order in which they occurred.
+        List<EventMetricData> data = getEventMetricDataList();
+
+        // Assert that the events happened in the expected order.
+        assertStatesOccurred(stateSet, data, WAIT_TIME_SHORT,
+            atom -> atom.getWifiLockStateChanged().getState().getNumber());
+    }
+
     public void testWifiScan() throws Exception {
         if (!TESTS_ENABLED) return;
         if (!hasFeature(FEATURE_WIFI, true)) return;