diff --git a/service/java/com/android/server/wifi/WifiLockManager.java b/service/java/com/android/server/wifi/WifiLockManager.java
index fe193f0..df9586f 100644
--- a/service/java/com/android/server/wifi/WifiLockManager.java
+++ b/service/java/com/android/server/wifi/WifiLockManager.java
@@ -71,7 +71,7 @@
         if (!isValidLockMode(lockMode)) {
             throw new IllegalArgumentException("lockMode =" + lockMode);
         }
-        if (ws == null || ws.size() == 0) {
+        if (ws == null || ws.isEmpty()) {
             ws = new WorkSource(Binder.getCallingUid());
         } else {
             mContext.enforceCallingOrSelfPermission(
@@ -144,15 +144,23 @@
         }
 
         WorkSource newWorkSource;
-        if (ws == null || ws.size() == 0) {
+        if (ws == null || ws.isEmpty()) {
             newWorkSource = new WorkSource(Binder.getCallingUid());
         } else {
             // Make a copy of the WorkSource before adding it to the WakeLock
             newWorkSource = new WorkSource(ws);
         }
 
+        if (mVerboseLoggingEnabled) {
+            Slog.d(TAG, "updateWifiLockWakeSource: " + wl + ", newWorkSource=" + newWorkSource);
+        }
+
         long ident = Binder.clearCallingIdentity();
         try {
+            // TODO: Introduce a noteFullWifiLockChangedFromSource() so that this logic can be
+            // handled properly in BatteryStats rather than here. This would ensure there are no
+            // "holes" in the reported data due to a "release" event occurring before an
+            // "acquire" event.
             mBatteryStats.noteFullWifiLockReleasedFromSource(wl.mWorkSource);
             wl.mWorkSource = newWorkSource;
             mBatteryStats.noteFullWifiLockAcquiredFromSource(wl.mWorkSource);
@@ -323,7 +331,8 @@
         }
 
         public String toString() {
-            return "WifiLock{" + this.mTag + " type=" + this.mMode + " uid=" + mUid + "}";
+            return "WifiLock{" + this.mTag + " type=" + this.mMode + " uid=" + mUid
+                    + " workSource=" + mWorkSource + "}";
         }
     }
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiLockManagerTest.java b/tests/wifitests/src/com/android/server/wifi/WifiLockManagerTest.java
index d5d79b8..1cc9479 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiLockManagerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiLockManagerTest.java
@@ -31,7 +31,6 @@
 
 import org.junit.Before;
 import org.junit.Test;
-
 import org.mockito.ArgumentCaptor;
 import org.mockito.InOrder;
 import org.mockito.Mock;
@@ -52,7 +51,9 @@
     WifiLockManager mWifiLockManager;
     @Mock IBatteryStats mBatteryStats;
     @Mock IBinder mBinder;
+    @Mock IBinder mBinder2;
     WorkSource mWorkSource;
+    WorkSource mChainedWorkSource;
     @Mock Context mContext;
 
     /**
@@ -62,6 +63,11 @@
     @Before
     public void setUp() {
         mWorkSource = new WorkSource(DEFAULT_TEST_UID_1);
+        mChainedWorkSource = new WorkSource();
+        mChainedWorkSource.createWorkChain()
+                .addNode(DEFAULT_TEST_UID_1, "tag1")
+                .addNode(DEFAULT_TEST_UID_2, "tag2");
+
         MockitoAnnotations.initMocks(this);
         mWifiLockManager = new WifiLockManager(mContext, mBatteryStats);
     }
@@ -225,6 +231,41 @@
     }
 
     /**
+     * Checks that WorkChains are preserved when merged WorkSources are created.
+     */
+    @Test
+    public void createMergedworkSourceWithChainsShouldSucceed() throws Exception {
+        acquireWifiLockSuccessful(WifiManager.WIFI_MODE_FULL_HIGH_PERF, "", mBinder, mWorkSource);
+        acquireWifiLockSuccessful(WifiManager.WIFI_MODE_FULL_HIGH_PERF, "", mBinder2,
+                mChainedWorkSource);
+
+        WorkSource merged = mWifiLockManager.createMergedWorkSource();
+        assertEquals(1, merged.size());
+        assertEquals(1, merged.getWorkChains().size());
+    }
+
+    /**
+     * A smoke test for acquiring, updating and releasing WifiLocks with chained WorkSources.
+     */
+    @Test
+    public void smokeTestLockLifecycleWithChainedWorkSource() throws Exception {
+        acquireWifiLockSuccessful(WifiManager.WIFI_MODE_FULL_HIGH_PERF, "", mBinder,
+                mChainedWorkSource);
+
+        WorkSource updated = new WorkSource();
+        updated.set(mChainedWorkSource);
+        updated.createWorkChain().addNode(
+                DEFAULT_TEST_UID_1, "chain2");
+
+        mWifiLockManager.updateWifiLockWorkSource(mBinder, updated);
+        InOrder inOrder = inOrder(mBatteryStats);
+        inOrder.verify(mBatteryStats).noteFullWifiLockReleasedFromSource(mChainedWorkSource);
+        inOrder.verify(mBatteryStats).noteFullWifiLockAcquiredFromSource(eq(updated));
+
+        releaseWifiLockSuccessful(mBinder);
+    }
+
+    /**
      * Test the ability to update a WifiLock WorkSource with a new WorkSource.
      *
      * Steps: acquire a WifiLock with the default test worksource, then attempt to update it.
@@ -314,6 +355,7 @@
         assertTrue(wifiLockManagerDumpString.contains("Locks held:"));
         assertTrue(wifiLockManagerDumpString.contains(
                 "WifiLock{" + TEST_WIFI_LOCK_TAG + " type=" + WifiManager.WIFI_MODE_FULL
-                + " uid=" + Binder.getCallingUid() + "}"));
+                + " uid=" + Binder.getCallingUid() + " workSource=WorkSource{"
+                        + DEFAULT_TEST_UID_1 + "}"));
     }
 }